def test_pyrocov_reparam(model, Guide, backend):
T, P, S, F = 2, 3, 4, 5
dataset = {
"features": torch.randn(S, F),
"local_time": torch.randn(T, P),
"weekly_strains": torch.randn(T, P, S).exp().round(),
}
# Reparametrize the model.
config = {
"coef": LocScaleReparam(),
"rate_loc": None if model is pyrocov_model else LocScaleReparam(),
"rate": LocScaleReparam(),
"init_loc": LocScaleReparam(),
"init": LocScaleReparam(),
}
model = poutine.reparam(model, config)
guide = Guide(model, backend=backend)
svi = SVI(model, guide, ClippedAdam({"lr": 1e-8}), Trace_ELBO())
for step in range(2):
with xfail_if_not_implemented():
svi.step(dataset)
guide(dataset)
predictive = Predictive(model, guide=guide, num_samples=2)
predictive(dataset)
def test_normal(dist_type, centered, shape):
loc = torch.empty(shape).uniform_(-1., 1.).requires_grad_()
scale = torch.empty(shape).uniform_(0.5, 1.5).requires_grad_()
if isinstance(centered, torch.Tensor):
centered = centered.expand(shape)
def model():
with pyro.plate_stack("plates", shape):
with pyro.plate("particles", 200000):
if "dist_type" == "Normal":
pyro.sample("x", dist.Normal(loc, scale))
else:
pyro.sample("x", dist.StudentT(10.0, loc, scale))
value = poutine.trace(model).get_trace().nodes["x"]["value"]
expected_probe = get_moments(value)
if "dist_type" == "Normal":
reparam = LocScaleReparam()
else:
reparam = LocScaleReparam(shape_params=["df"])
reparam_model = poutine.reparam(model, {"x": reparam})
value = poutine.trace(reparam_model).get_trace().nodes["x"]["value"]
actual_probe = get_moments(value)
assert_close(actual_probe, expected_probe, atol=0.1)
for actual_m, expected_m in zip(actual_probe, expected_probe):
expected_grads = grad(expected_m.sum(), [loc, scale],
retain_graph=True)
actual_grads = grad(actual_m.sum(), [loc, scale], retain_graph=True)
assert_close(actual_grads[0], expected_grads[0], atol=0.05)
assert_close(actual_grads[1], expected_grads[1], atol=0.05)
def model(self, zero_data, covariates):
data_dim = zero_data.size(-1)
feature_dim = covariates.size(-1)
bias = pyro.sample("bias", dist.Normal(5.0, 10.0).expand((data_dim,)).to_event(1))
weight = pyro.sample("weight", dist.Normal(0.0, 5).expand((feature_dim,)).to_event(1))
# We'll sample a time-global scale parameter outside the time plate,
# then time-local iid noise inside the time plate.
drift_scale = pyro.sample("drift_scale",
dist.LogNormal(-20, 5.0).expand((1,)).to_event(1))
with self.time_plate:
# We'll use a reparameterizer to improve variational fit. The model would still be
# correct if you removed this context manager, but the fit appears to be worse.
with poutine.reparam(config={"drift": LocScaleReparam()}):
drift = pyro.sample("drift", dist.Normal(zero_data.double(), drift_scale.double()).to_event(1))
# After we sample the iid "drift" noise we can combine it in any time-dependent way.
# It is important to keep everything inside the plate independent and apply dependent
# transforms outside the plate.
motion = drift.cumsum(-2) # A Brownian motion.
# The prediction now includes three terms.
prediction = motion + bias + (weight * covariates).sum(-1, keepdim=True)
assert prediction.shape[-2:] == zero_data.shape
# Construct the noise distribution and predict.
noise_scale = pyro.sample("noise_scale", dist.LogNormal(0.0, 1.0).expand((1,)).to_event(1))
noise_dist = dist.Normal(0, noise_scale)
self.predict(noise_dist, prediction)
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_init(dist_type, centered, shape):
loc = torch.empty(shape).uniform_(-1.0, 1.0)
scale = torch.empty(shape).uniform_(0.5, 1.5)
def model():
with pyro.plate_stack("plates", shape):
if "dist_type" == "Normal":
return pyro.sample("x", dist.Normal(loc, scale))
elif "dist_type" == "StudentT":
return pyro.sample("x", dist.StudentT(10.0, loc, scale))
else:
return pyro.sample("x",
dist.AsymmetricLaplace(loc, scale, 1.5))
check_init_reparam(model, LocScaleReparam(centered))
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)
def model(self, zero_data, covariates):
data_dim = zero_data.size(-1)
feature_dim = covariates.size(-1)
bias = pyro.sample("bias", dist.Normal(5.0, 10.0).expand((data_dim,)).to_event(1))
weight = pyro.sample("weight", dist.Normal(0.0, 5).expand((feature_dim,)).to_event(1))
# We'll sample a time-global scale parameter outside the time plate,
# then time-local iid noise inside the time plate.
drift_scale = pyro.sample("drift_scale",
dist.LogNormal(0, 5.0).expand((1,)).to_event(1))
with self.time_plate:
# We'll use a reparameterizer to improve variational fit. The model would still be
# correct if you removed this context manager, but the fit appears to be worse.
with poutine.reparam(config={"drift": LocScaleReparam()}):
drift = pyro.sample("drift", dist.Normal(zero_data.double(), drift_scale.double()).to_event(1))
# After we sample the iid "drift" noise we can combine it in any time-dependent way.
# It is important to keep everything inside the plate independent and apply dependent
# transforms outside the plate.
motion = drift.cumsum(-2) # A Brownian motion.
# The prediction now includes three terms.
prediction = motion + bias + (weight * covariates).sum(-1, keepdim=True)
assert prediction.shape[-2:] == zero_data.shape
# The next part of the model creates a likelihood or noise distribution.
# Again we'll be Bayesian and write this as a probabilistic program with
# priors over parameters.
stability = pyro.sample("noise_stability", dist.Uniform(1, 2).expand((1,)).to_event(1))
skew = pyro.sample("noise_skew", dist.Uniform(-1, 1).expand((1,)).to_event(1))
scale = pyro.sample("noise_scale", dist.LogNormal(-5, 5).expand((1,)).to_event(1))
noise_dist = dist.Stable(stability, skew, scale)
# We need to use a reparameterizer to handle the Stable distribution.
# Note "residual" is the name of Pyro's internal sample site in self.predict().
with poutine.reparam(config={"residual": StableReparam()}):
self.predict(noise_dist, prediction)
def model():
with pyro.plate("N", 10):
with poutine.reparam(config={"x": LocScaleReparam(centered=0.0)}):
return pyro.sample("x", dist.Normal(0, 1))