def model(self, zero_data, covariates):
with pyro.plate("batch", len(zero_data), dim=-2):
zero_data = pyro.subsample(zero_data, event_dim=1)
covariates = pyro.subsample(covariates, event_dim=1)
loc = zero_data[..., :1, :]
scale = pyro.sample("scale", dist.LogNormal(loc, 1).to_event(1))
with self.time_plate:
jumps = pyro.sample("jumps", dist.Normal(0, scale).to_event(1))
prediction = jumps.cumsum(-2)
duration, obs_dim = zero_data.shape[-2:]
noise_dist = dist.LinearHMM(
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
duration=duration,
)
rep = StableReparam()
with poutine.reparam(
config={"residual": LinearHMMReparam(rep, rep, rep)}):
self.predict(noise_dist, prediction)
def model(self, zero_data, covariates):
assert zero_data.size(-1) == 1 # univariate
num_stores, num_products, duration, one = zero_data.shape
time_index = covariates.squeeze(-1)
store_plate = pyro.plate("store", num_stores, dim=-3)
product_plate = pyro.plate("product", num_products, dim=-2)
day_of_week_plate = pyro.plate("day_of_week", 7, dim=-1)
snap = self.snap[..., time_index, :]
# subsample the data
with product_plate:
dept = pyro.subsample(self.dept, event_dim=1)
saled = pyro.subsample(self.saled, event_dim=1)[..., time_index, :]
log_ma = pyro.subsample(self.log_ma, event_dim=1)[...,
time_index, :]
# we construct latent variables for each store and each department;
# here, we declare department dimension as event dimension for simplicity,
# (nb: the numbers of products in each department are different)
# the last event dimension is used to model mean/scale separately.
with store_plate:
ma_weight = pyro.sample(
"ma_weight",
dist.Normal(0, 1).expand([2, log_ma.size(-1), 7]).to_event(3))
ma_weight = ma_weight.matmul(
dept.unsqueeze(-2).unsqueeze(-1)).squeeze(-1)
moving_average = ma_weight.matmul(log_ma.unsqueeze(-1)).squeeze(-1)
snap_weight = pyro.sample(
"snap_weight",
dist.Normal(0, 1).expand([2, 7]).to_event(2))
snap_weight = snap_weight.matmul(dept.unsqueeze(-1)).squeeze(-1)
snap_effect = snap_weight * snap
with day_of_week_plate:
seasonal = pyro.sample(
"seasonal",
dist.Normal(0, 1).expand([2, 7]).to_event(2))
seasonal = seasonal.matmul(dept.unsqueeze(-1)).squeeze(-1)
seasonal = periodic_repeat(seasonal, duration, dim=-2)
prediction = moving_average + snap_effect + seasonal
log_mean, log_scale = prediction[..., :1], prediction[..., 1:]
# we add a pretty small bias 1e-3 to avoid the case mean=scale=0
# either when saled == 0 or saled == 1
mean = bounded_exp(log_mean) * saled + 1e-3
scale = bounded_exp(log_scale) * saled + 1e-3
rate = scale.reciprocal()
concentration = mean * rate
# alternative: GammaPoisson (or NegativeBinomial, ZeroInflatedNegativeBinomial)
noise_dist = dist.Gamma(concentration, rate)
with store_plate, product_plate:
self.predict(noise_dist, mean.new_zeros(mean.shape))
def model(x, y=None, batch_size=None):
loc = pyro.param("loc", lambda: torch.tensor(0.))
scale = pyro.param("scale", lambda: torch.tensor(1.),
constraint=constraints.positive)
with pyro.plate("batch", len(x), subsample_size=batch_size):
batch_x = pyro.subsample(x, event_dim=0)
batch_y = pyro.subsample(y, event_dim=0) if y is not None else None
mean = loc + scale * batch_x
sigma = pyro.sample("sigma", dist.LogNormal(0., 1.))
return pyro.sample("obs", dist.Normal(mean, sigma), obs=batch_y)
def model(data):
size, size = data.shape
origin_plate = pyro.plate("origin", size, dim=-2)
destin_plate = pyro.plate("destin", size, dim=-1)
with origin_plate, destin_plate:
batch = pyro.subsample(data, event_dim=0)
assert batch.size(0) == batch.size(1), batch.shape
pyro.sample("obs", dist.Normal(0, 1), obs=batch)
def model(self, zero_data, covariates):
with pyro.plate("batch", len(zero_data), dim=-2):
zero_data = pyro.subsample(zero_data, event_dim=1)
covariates = pyro.subsample(covariates, event_dim=1)
loc = zero_data[..., :1, :]
scale = pyro.sample("scale", dist.LogNormal(loc, 1).to_event(1))
with self.time_plate:
jumps = pyro.sample("jumps", dist.Normal(0, scale).to_event(1))
prediction = jumps.cumsum(-2)
duration, obs_dim = zero_data.shape[-2:]
noise_dist = dist.GaussianHMM(
dist.Normal(0, 1).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Normal(0, 1).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Normal(0, 1).expand([obs_dim]).to_event(1),
duration=duration,
)
self.predict(noise_dist, prediction)
def predict(self, noise_dist, prediction):
"""
Prediction function, to be called by :meth:`model` implementations.
This should be called outside of the :meth:`time_plate`.
This is similar to an observe statement in Pyro::
pyro.sample("residual", noise_dist,
obs=(data - prediction))
but with (1) additional reshaping logic to allow time-dependent
``noise_dist`` (most often a :class:`~pyro.distributions.GaussianHMM`
or variant); and (2) additional logic to allow only a partial
observation and forecast the remaining data.
:param noise_dist: A noise distribution with ``.event_dim in {0,1,2}``.
``noise_dist`` is typically zero-mean or zero-median or zero-mode
or somehow centered.
:type noise_dist: ~pyro.distributions.Distribution
:param prediction: A prediction for the data. This should have the same
shape as ``data``, but broadcastable to full duration of the
``covariates``.
:type prediction: ~torch.Tensor
"""
assert self._data is not None, ".predict() called outside .model()"
assert self._forecast is None, ".predict() called twice"
assert isinstance(noise_dist, dist.Distribution)
assert isinstance(prediction, torch.Tensor)
if noise_dist.event_dim == 0:
if noise_dist.batch_shape[-2:] != prediction.shape[-2:]:
noise_dist = noise_dist.expand(noise_dist.batch_shape[:-2] +
prediction.shape[-2:])
noise_dist = noise_dist.to_event(2)
elif noise_dist.event_dim == 1:
if noise_dist.batch_shape[-1:] != prediction.shape[-2:-1]:
noise_dist = noise_dist.expand(noise_dist.batch_shape[:-1] +
prediction.shape[-2:-1])
noise_dist = noise_dist.to_event(1)
assert noise_dist.event_dim == 2
assert noise_dist.event_shape == prediction.shape[-2:]
# The following reshaping logic is required to reconcile batch and
# event shapes. This would be unnecessary if Pyro used name dimensions
# internally, e.g. using Funsor.
#
# batch_shape | event_shape
# -------------------------------+----------------
# 1. sample_shape + shape + (time,) | (obs_dim,)
# 2. sample_shape + shape | (time, obs_dim)
# 3. sample_shape + shape + (1,) | (time, obs_dim)
#
# Parameters like noise_dist.loc typically have shape as in 1. However
# calling .to_event(1) will shift the shapes resulting in 2., where
# sample_shape+shape will be misaligned with other batch shapes in the
# trace. To fix this the following logic "unsqueezes" the distribution,
# resulting in correctly aligned shapes 3. Note the "time" dimension is
# effectively moved from a batch dimension to an event dimension.
noise_dist = reshape_batch(noise_dist, noise_dist.batch_shape + (1, ))
data = pyro.subsample(self._data.unsqueeze(-3), event_dim=2)
prediction = prediction.unsqueeze(-3)
# Create a sample site.
t_obs = data.size(-2)
t_cov = prediction.size(-2)
if t_obs == t_cov: # training
pyro.sample("residual", noise_dist, obs=data - prediction)
self._forecast = data.new_zeros(data.shape[:-2] + (0, ) +
data.shape[-1:])
else: # forecasting
left_pred = prediction[..., :t_obs, :]
right_pred = prediction[..., t_obs:, :]
# This prefix_condition indirection is needed to ensure that
# PrefixConditionMessenger is handled outside of the .model() call.
self._prefix_condition_data["residual"] = data - left_pred
noise = pyro.sample("residual", noise_dist)
del self._prefix_condition_data["residual"]
assert noise.shape[-data.dim():] == right_pred.shape[-data.dim():]
self._forecast = right_pred + noise
# Move the "time" batch dim back to its original place.
assert self._forecast.size(-3) == 1
self._forecast = self._forecast.squeeze(-3)