def forward(self, features, trip_counts):
pyro.module("model", self)
total_hours = len(features)
observed_hours, num_origins, num_destins = trip_counts.shape
assert observed_hours <= total_hours
assert num_origins == self.num_stations
assert num_destins == self.num_stations
time_plate = pyro.plate("time", observed_hours, dim=-3)
origins_plate = pyro.plate("origins", num_origins, dim=-2)
destins_plate = pyro.plate("destins", num_destins, dim=-1)
# The first half of the model performs exact inference over
# the observed portion of the time series.
hmm = dist.GaussianHMM(*self._dynamics(features[:observed_hours]))
gate_rate = pyro.sample("gate_rate", hmm)
gate, rate = self._unpack_gate_rate(gate_rate, event_dim=2)
with time_plate, origins_plate, destins_plate:
pyro.sample("trip_count",
dist.ZeroInflatedPoisson(gate, rate),
obs=trip_counts)
# The second half of the model forecasts forward.
forecast = []
forecast_hours = total_hours - observed_hours
if forecast_hours > 0:
_, trans_matrix, trans_dist, obs_matrix, obs_dist = \
self._dynamics(features[observed_hours:])
state = None
for t in range(forecast_hours):
if state is None: # on first step
state_dist = hmm.filter(gate_rate)
else:
loc = vm(state, trans_matrix) + trans_dist.loc
scale_tril = trans_dist.scale_tril
state_dist = dist.MultivariateNormal(loc,
scale_tril=scale_tril)
state = pyro.sample("state_{}".format(t), state_dist)
loc = vm(state, obs_matrix) + obs_dist.base_dist.loc[..., t, :]
scale = obs_dist.base_dist.scale[..., t, :]
gate_rate = pyro.sample("gate_rate_{}".format(t),
dist.Normal(loc, scale).to_event(1))
gate, rate = self._unpack_gate_rate(gate_rate, event_dim=1)
with origins_plate, destins_plate:
forecast.append(
pyro.sample("trip_count_{}".format(t),
dist.ZeroInflatedPoisson(gate, rate)))
return forecast
def _forward_pyro_forecast(self,
features,
trip_counts,
origins_plate,
destins_plate,
state=None,
state_dist=None):
total_hours = len(features)
observed_hours, num_origins, num_destins = trip_counts.shape
forecast = []
forecast_hours = total_hours - observed_hours
if forecast_hours > 0:
_, trans_matrix, trans_dist, obs_matrix, obs_dist = \
self._dynamics(features[observed_hours:])
for t in range(forecast_hours):
if state is not None:
loc = vm(state, trans_matrix) + trans_dist.loc
scale_tril = trans_dist.scale_tril
state_dist = dist.MultivariateNormal(loc,
scale_tril=scale_tril)
state = pyro.sample("state_{}".format(t), state_dist)
loc = vm(state, obs_matrix) + obs_dist.base_dist.loc[..., t, :]
scale = obs_dist.base_dist.scale[..., t, :]
gate_rate = pyro.sample("gate_rate_{}".format(t),
dist.Normal(loc, scale).to_event(1))
gate, rate = self._unpack_gate_rate(gate_rate, event_dim=1)
with origins_plate, destins_plate:
forecast.append(
pyro.sample("trip_count_{}".format(t),
dist.ZeroInflatedPoisson(gate, rate)))
return forecast
def _forward_pyro(self, features, trip_counts):
total_hours = len(features)
observed_hours, num_origins, num_destins = trip_counts.shape
assert observed_hours <= total_hours
assert num_origins == self.num_stations
assert num_destins == self.num_stations
time_plate = pyro.plate("time", observed_hours, dim=-3)
origins_plate = pyro.plate("origins", num_origins, dim=-2)
destins_plate = pyro.plate("destins", num_destins, dim=-1)
# The first half of the model performs exact inference over
# the observed portion of the time series.
hmm = dist.GaussianHMM(*self._dynamics(features[:observed_hours]))
gate_rate = pyro.sample("gate_rate", hmm)
gate, rate = self._unpack_gate_rate(gate_rate, event_dim=2)
with time_plate, origins_plate, destins_plate:
pyro.sample("trip_count",
dist.ZeroInflatedPoisson(gate, rate),
obs=trip_counts)
# The second half of the model forecasts forward.
if total_hours > observed_hours:
state_dist = hmm.filter(gate_rate)
return self._forward_pyro_forecast(features,
trip_counts,
origins_plate,
destins_plate,
state_dist=state_dist)
def test_zip_shape():
gate = torch.ones(3, 2) / 2
rate = torch.ones(3, 2) / 2
d = dist.ZeroInflatedPoisson(rate, gate=gate)
assert d.batch_shape == (3, 2)
assert d.event_shape == ()
assert d.shape() == (3, 2)
assert d.sample().size() == d.shape()
def post_model(
p_data,
t_data,
s_data,
r_data,
y,
p_types,
p_stories,
p_subreddits,
zero_inflated,
):
coef_scale_prior = 0.1
num_posts, num_p_indeps = p_data.shape
# shared prior
gamma_loc = torch.zeros((num_p_indeps, 1), dtype=torch.float64)
if zero_inflated:
gamma_gate_loc = torch.zeros((num_p_indeps, 1), dtype=torch.float64)
with pyro.plate("p_indep", num_p_indeps, dim=-2):
gamma = pyro.sample("gamma", dist.Normal(gamma_loc, coef_scale_prior))
if zero_inflated:
gamma_gate = pyro.sample(
"gamma_gate", dist.Normal(gamma_gate_loc, coef_scale_prior)
)
# for each post,
# use the correct set of coefficients to run our post-level regression
with pyro.plate("post", num_posts, dim=-1) as p:
# indep vars for this post
indeps = p_data[p, :]
mu = torch.matmul(
indeps, gamma
).flatten() # ( num_posts, num_p_indeps) x (num_p_indeps, 1)
# defining response dist
if zero_inflated:
gate = torch.nn.Sigmoid()(
torch.matmul(indeps, gamma_gate).flatten()
) # ( num_posts, num_p_indeps) x (num_p_indeps, 1)
response_dist = dist.ZeroInflatedPoisson(
rate=torch.exp(mu), gate=gate
)
else:
response_dist = dist.Poisson(rate=torch.exp(mu))
# sample
if y is None:
pyro.sample("obs", response_dist, obs=y)
else:
pyro.sample("obs", response_dist, obs=y[p])
def _forward_pyro_mean_field(self, features, trip_counts):
total_hours = len(features)
observed_hours, num_origins, num_destins = trip_counts.shape
assert observed_hours <= total_hours
assert num_origins == self.num_stations
assert num_destins == self.num_stations
time_plate = pyro.plate("time", observed_hours, dim=-3)
origins_plate = pyro.plate("origins", num_origins, dim=-2)
destins_plate = pyro.plate("destins", num_destins, dim=-1)
init_dist, trans_matrix, trans_dist, obs_matrix, obs_dist = \
self._dynamics(features[:observed_hours])
# This is a parallelizable crf representation of the HMM.
# We first pull random variables from the guide, masking all factors.
with poutine.mask(mask=False):
shape = (1 + observed_hours, self.args.state_dim) # includes init
state = pyro.sample("state",
dist.Normal(0, 1).expand(shape).to_event(2))
shape = (observed_hours, 2 * num_origins * num_destins)
gate_rate = pyro.sample(
"gate_rate",
dist.Normal(0, 1).expand(shape).to_event(2))
# We then declare CRF factors.
pyro.sample("init", init_dist, obs=state[0])
pyro.sample("trans",
trans_dist.expand((observed_hours, )).to_event(1),
obs=state[..., 1:, :] - state[..., :-1, :] @ trans_matrix)
pyro.sample("obs",
obs_dist.expand((observed_hours, )).to_event(1),
obs=gate_rate - state[..., 1:, :] @ obs_matrix)
gate, rate = self._unpack_gate_rate(gate_rate, event_dim=2)
with time_plate, origins_plate, destins_plate:
pyro.sample("trip_count",
dist.ZeroInflatedPoisson(gate, rate),
obs=trip_counts)
# The second half of the model forecasts forward.
if total_hours > observed_hours:
return self._forward_pyro_forecast(features,
trip_counts,
origins_plate,
destins_plate,
state=state[..., -1, :])
def model(self, data, demand):
coef = {}
for s in self.features['station']['names']:
coef[s] = pyro.sample(s, dist.Normal(0, 1))
s += '_gate'
coef[s] = pyro.sample(s, dist.Normal(0, 1))
for h in self.features['hour']['names']:
for d in self.features['daytype']['names']:
name = h + '_' + d
coef[name] = pyro.sample(name, dist.Normal(0, 1))
name += '_gate'
coef[name] = pyro.sample(name, dist.Normal(0, 1))
log_lmbda = 0
gate_mean = 0
for i in range(len(self.features['station']['names'])):
name = self.features['station']['names'][i]
index = self.features['station']['index'][i]
log_lmbda += coef[name] * data[:, index]
gate_mean += coef[name + '_gate'] * data[:, index]
for h in range(len(self.features['hour']['names'])):
for d in range(len(self.features['daytype']['names'])):
h_name = self.features['hour']['names'][h]
h_index = self.features['hour']['index'][h]
d_name = self.features['daytype']['names'][d]
d_index = self.features['daytype']['index'][d]
log_lmbda += coef[h_name + '_' + d_name] * \
data[:, h_index] * data[:, d_index]
gate_mean += coef[h_name + '_' + d_name + '_gate'] * \
data[:, h_index] * data[:, d_index]
lmbda = log_lmbda.exp()
gate = sigmoid(gate_mean)
with pyro.plate("data", len(data)):
pyro.sample("obs",
dist.ZeroInflatedPoisson(gate, lmbda),
obs=demand)
return gate, lmbda
def model(self, data, demand):
coef = {}
for s in self.features['station']['names']:
coef[s] = pyro.sample(s, dist.Normal(0, 1))
for h in self.features['hour']['names']:
for d in self.features['daytype']['names']:
name = h + '_' + d
coef[name] = pyro.sample(name, dist.Normal(0, 1))
log_lmbda = 0
for i in range(len(self.features['station']['names'])):
name = self.features['station']['names'][i]
index = self.features['station']['index'][i]
log_lmbda += coef[name] * data[:, index]
for h in range(len(self.features['hour']['names'])):
for d in range(len(self.features['daytype']['names'])):
h_name = self.features['hour']['names'][h]
h_index = self.features['hour']['index'][h]
d_name = self.features['daytype']['names'][d]
d_index = self.features['daytype']['index'][d]
log_lmbda += coef[h_name + '_' + d_name] * \
data[:, h_index] * data[:, d_index]
lmbda = log_lmbda.exp()
gate_alpha = pyro.sample('gate_alpha', dist.Gamma(2, 2))
gate_beta = pyro.sample('gate_beta', dist.Gamma(3, 2))
gate = pyro.sample('gate', dist.Beta(gate_alpha, gate_beta))
with pyro.plate("data", len(data)):
pyro.sample(
"obs", dist.ZeroInflatedPoisson(
gate, lmbda), obs=demand)
# should we be returning lmbda?
return gate, lmbda
def type_model(
p_data,
t_data,
s_data,
r_data,
y,
p_types,
p_stories,
p_subreddits,
zero_inflated,
):
coef_scale_prior = 0.1
num_posts, num_p_indeps = p_data.shape
num_types, num_t_indeps = t_data.shape
# type priors
alpha_loc = torch.zeros((num_p_indeps, num_t_indeps), dtype=torch.float64)
alpha_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_t_indeps), dtype=torch.float64)
if zero_inflated:
eta_gate_loc = torch.zeros((num_p_indeps, num_t_indeps),
dtype=torch.float64)
eta_gate_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_t_indeps), dtype=torch.float64)
with pyro.plate("p_indep", num_p_indeps, dim=-2):
# Type Level
with pyro.plate("t_indep", num_t_indeps, dim=-1):
eta = pyro.sample("eta", dist.Normal(alpha_loc, alpha_scale))
if zero_inflated:
eta_gate = pyro.sample(
"eta_gate", dist.Normal(eta_gate_loc, eta_gate_scale))
with pyro.plate("type", num_types, dim=-1) as t:
phi_loc = torch.matmul(
eta, t_data[t, :].T
) # (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
phi = pyro.sample("phi", dist.Normal(phi_loc, coef_scale_prior))
if zero_inflated:
phi_gate_loc = torch.matmul(
eta_gate, t_data[t, :].T
) # (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
phi_gate = pyro.sample(
"phi_gate", dist.Normal(phi_gate_loc, coef_scale_prior))
# for each post,
# use the correct set of coefficients to run our post-level regression
with pyro.plate("post", num_posts, dim=-1) as p:
t = p_types[p]
# indep vars for this post
indeps = p_data[p, :]
t_coefs = phi[:, t] # (num_p_indeps,num_posts)
type_level_products = torch.mul(
t_coefs,
indeps.T) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
# calculate the mean: desired shape (num_posts, 1)
mu = (type_level_products).sum(
dim=0) # (num_p_indeps, num_posts).sum(over indeps)
# defining response dist
if zero_inflated:
t_coefs_gate = phi_gate[:, t] # (num_p_indeps,num_posts)
type_level_products_gate = torch.mul(
t_coefs_gate, indeps.T
) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
# calculate the mean: desired shape (num_posts, 1)
gate = torch.nn.Sigmoid()((type_level_products_gate).sum(
dim=0)) # (num_p_indeps, num_posts).sum(over indeps)
response_dist = dist.ZeroInflatedPoisson(rate=torch.exp(mu),
gate=gate)
else:
response_dist = dist.Poisson(rate=torch.exp(mu))
# sample
if y is None:
pyro.sample("obs", response_dist, obs=y)
else:
pyro.sample("obs", response_dist, obs=y[p])
def complete_model(
p_data,
t_data,
s_data,
r_data,
y,
p_types,
p_stories,
p_subreddits,
zero_inflated,
):
coef_scale_prior = 0.1
num_posts, num_p_indeps = p_data.shape
num_types, num_t_indeps = t_data.shape
num_stories, num_s_indeps = s_data.shape
num_subreddits, num_r_indeps = r_data.shape
# type priors
alpha_loc = torch.zeros((num_p_indeps, num_t_indeps), dtype=torch.float64)
alpha_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_t_indeps), dtype=torch.float64)
# story priors
beta_loc = torch.zeros((num_p_indeps, num_s_indeps), dtype=torch.float64)
beta_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_s_indeps), dtype=torch.float64)
# subreddit priors
tau_loc = torch.zeros((num_p_indeps, num_r_indeps), dtype=torch.float64)
tau_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_r_indeps), dtype=torch.float64)
if zero_inflated:
# type priors
eta_gate_loc = torch.zeros((num_p_indeps, num_t_indeps),
dtype=torch.float64)
eta_gate_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_t_indeps), dtype=torch.float64)
# story priors
beta_gate_loc = torch.zeros((num_p_indeps, num_s_indeps),
dtype=torch.float64)
beta_gate_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_s_indeps), dtype=torch.float64)
# subreddit priors
tau_gate_loc = torch.zeros((num_p_indeps, num_r_indeps),
dtype=torch.float64)
tau_gate_scale = coef_scale_prior * torch.ones(
(num_p_indeps, num_r_indeps), dtype=torch.float64)
with pyro.plate("p_indep", num_p_indeps, dim=-2):
# Type Level
with pyro.plate("t_indep", num_t_indeps, dim=-1):
eta = pyro.sample("eta", dist.Normal(alpha_loc, alpha_scale))
if zero_inflated:
eta_gate = pyro.sample(
"eta_gate", dist.Normal(eta_gate_loc, eta_gate_scale))
with pyro.plate("type", num_types, dim=-1) as t:
phi_loc = torch.matmul(eta, t_data[t, :].T)
# (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
phi = pyro.sample("phi", dist.Normal(phi_loc, coef_scale_prior))
if zero_inflated:
phi_gate_loc = torch.matmul(
eta_gate, t_data[t, :].T
) # (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
phi_gate = pyro.sample(
"phi_gate", dist.Normal(phi_gate_loc, coef_scale_prior))
# Story Level
with pyro.plate("s_indep", num_s_indeps, dim=-1):
beta = pyro.sample("beta", dist.Normal(beta_loc, beta_scale))
if zero_inflated:
beta_gate = pyro.sample(
"beta_gate", dist.Normal(beta_gate_loc, beta_gate_scale))
with pyro.plate("story", num_stories, dim=-1) as s:
theta_loc = torch.matmul(
beta, s_data[s, :].T
) # (num_p_indeps, num_s_indeps) x (num_s_indeps, num_stories)
theta = pyro.sample("theta",
dist.Normal(theta_loc, coef_scale_prior))
if zero_inflated:
theta_gate_loc = torch.matmul(
beta_gate, s_data[s, :].T
) # (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
theta_gate = pyro.sample(
"theta_gate", dist.Normal(theta_gate_loc,
coef_scale_prior))
# Subreddit Level
with pyro.plate("r_indep", num_r_indeps, dim=-1):
tau = pyro.sample("tau", dist.Normal(tau_loc, tau_scale))
if zero_inflated:
tau_gate = pyro.sample(
"tau_gate", dist.Normal(tau_gate_loc, tau_gate_scale))
with pyro.plate("subreddit", num_subreddits, dim=-1) as r:
rho_loc = torch.matmul(
tau, r_data[r, :].T
) # (num_p_indeps, num_r_indeps) x (num_r_indeps, num_subreddits)
rho = pyro.sample("rho", dist.Normal(rho_loc, coef_scale_prior))
if zero_inflated:
rho_gate_loc = torch.matmul(
tau_gate, r_data[r, :].T
) # (num_p_indeps, num_t_indeps) x (num_t_indeps, num_types)
rho_gate = pyro.sample(
"rho_gate", dist.Normal(rho_gate_loc, coef_scale_prior))
# for each post,
# use the correct set of coefficients to run our post-level regression
with pyro.plate("post", num_posts, dim=-1) as p:
t = p_types[p]
s = p_stories[p]
r = p_subreddits[p]
# indep vars for this post
indeps = p_data[p, :]
t_coefs = phi[:, t] # (num_p_indeps,num_posts)
s_coefs = theta[:, s] # (num_p_indeps,num_posts)
r_coefs = rho[:, r] # (num_p_indeps,num_posts)
type_level_products = torch.mul(
t_coefs,
indeps.T) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
story_level_products = torch.mul(
s_coefs,
indeps.T) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
subreddit_level_products = torch.mul(
r_coefs,
indeps.T) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
# calculate the mean: desired shape (num_posts, 1)
mu = (subreddit_level_products + type_level_products +
story_level_products).sum(
dim=0) # (num_p_indeps, num_posts).sum(over indeps)
# defining response dist
if zero_inflated:
t_coefs_gate = phi_gate[:, t] # (num_p_indeps,num_posts)
s_coefs_gate = theta_gate[:, s] # (num_p_indeps,num_posts)
r_coefs_gate = rho_gate[:, r] # (num_p_indeps,num_posts)
type_level_products_gate = torch.mul(
t_coefs_gate, indeps.T
) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
story_level_products_gate = torch.mul(
s_coefs_gate, indeps.T
) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
subreddit_level_products_gate = torch.mul(
r_coefs_gate, indeps.T
) # (num_p_indeps, num_posts) .* (num_p_indeps, num_posts)
# calculate the mean: desired shape (num_posts, 1)
gate = torch.nn.Sigmoid()(
(type_level_products_gate + story_level_products_gate +
subreddit_level_products_gate).sum(
dim=0)) # (num_p_indeps, num_posts).sum(over indeps)
response_dist = dist.ZeroInflatedPoisson(rate=torch.exp(mu),
gate=gate)
else:
response_dist = dist.Poisson(rate=torch.exp(mu))
# sample
if y is None:
pyro.sample("obs", response_dist, obs=y)
else:
pyro.sample("obs", response_dist, obs=y[p])