def test_gamma_gaussian_hmm_log_prob(sample_shape, batch_shape, num_steps,
hidden_dim, obs_dim):
init_dist = random_mvn(batch_shape, hidden_dim)
trans_mat = torch.randn(batch_shape + (num_steps, hidden_dim, hidden_dim))
trans_dist = random_mvn(batch_shape + (num_steps, ), hidden_dim)
obs_mat = torch.randn(batch_shape + (num_steps, hidden_dim, obs_dim))
obs_dist = random_mvn(batch_shape + (num_steps, ), obs_dim)
scale_dist = random_gamma(batch_shape)
d = dist.GammaGaussianHMM(scale_dist, init_dist, trans_mat, trans_dist,
obs_mat, obs_dist)
obs_mvn = obs_dist
data = obs_dist.sample(sample_shape)
assert data.shape == sample_shape + d.shape()
actual_log_prob = d.log_prob(data)
# Compare against hand-computed density.
# We will construct enormous unrolled joint gaussian-gammas with shapes:
# t | 0 1 2 3 1 2 3 T = 3 in this example
# ------+-----------------------------------------
# init | H
# trans | H H H H H = hidden
# obs | H H H O O O O = observed
# and then combine these using gamma_gaussian_tensordot().
T = num_steps
init = gamma_and_mvn_to_gamma_gaussian(scale_dist, init_dist)
trans = matrix_and_mvn_to_gamma_gaussian(trans_mat, trans_dist)
obs = matrix_and_mvn_to_gamma_gaussian(obs_mat, obs_mvn)
unrolled_trans = reduce(
operator.add,
[
trans[..., t].event_pad(left=t * hidden_dim,
right=(T - t - 1) * hidden_dim)
for t in range(T)
],
)
unrolled_obs = reduce(
operator.add,
[
obs[..., t].event_pad(left=t * obs.dim(),
right=(T - t - 1) * obs.dim())
for t in range(T)
],
)
# Permute obs from HOHOHO to HHHOOO.
perm = torch.cat(
[torch.arange(hidden_dim) + t * obs.dim() for t in range(T)] +
[torch.arange(obs_dim) + hidden_dim + t * obs.dim() for t in range(T)])
unrolled_obs = unrolled_obs.event_permute(perm)
unrolled_data = data.reshape(data.shape[:-2] + (T * obs_dim, ))
assert init.dim() == hidden_dim
assert unrolled_trans.dim() == (1 + T) * hidden_dim
assert unrolled_obs.dim() == T * (hidden_dim + obs_dim)
logp = gamma_gaussian_tensordot(init, unrolled_trans, hidden_dim)
logp = gamma_gaussian_tensordot(logp, unrolled_obs, T * hidden_dim)
# compute log_prob of the joint student-t distribution
expected_log_prob = logp.compound().log_prob(unrolled_data)
assert_close(actual_log_prob, expected_log_prob)
def __init__(self,
scale_dist,
initial_dist,
transition_matrix,
transition_dist,
observation_matrix,
observation_dist,
validate_args=None):
assert isinstance(scale_dist, Gamma)
assert isinstance(initial_dist, MultivariateNormal)
assert isinstance(transition_matrix, torch.Tensor)
assert isinstance(transition_dist, MultivariateNormal)
assert isinstance(observation_matrix, torch.Tensor)
assert isinstance(observation_dist, MultivariateNormal)
hidden_dim, obs_dim = observation_matrix.shape[-2:]
assert initial_dist.event_shape == (hidden_dim, )
assert transition_matrix.shape[-2:] == (hidden_dim, hidden_dim)
assert transition_dist.event_shape == (hidden_dim, )
assert observation_dist.event_shape == (obs_dim, )
shape = broadcast_shape(scale_dist.batch_shape + (1, ),
initial_dist.batch_shape + (1, ),
transition_matrix.shape[:-2],
transition_dist.batch_shape,
observation_matrix.shape[:-2],
observation_dist.batch_shape)
batch_shape, time_shape = shape[:-1], shape[-1:]
event_shape = time_shape + (obs_dim, )
super(GammaGaussianHMM, self).__init__(batch_shape,
event_shape,
validate_args=validate_args)
self.hidden_dim = hidden_dim
self.obs_dim = obs_dim
self._init = gamma_and_mvn_to_gamma_gaussian(scale_dist, initial_dist)
self._trans = matrix_and_mvn_to_gamma_gaussian(transition_matrix,
transition_dist)
self._obs = matrix_and_mvn_to_gamma_gaussian(observation_matrix,
observation_dist)
def test_matrix_and_mvn_to_gamma_gaussian(sample_shape, batch_shape, x_dim,
y_dim):
matrix = torch.randn(batch_shape + (x_dim, y_dim))
y_mvn = random_mvn(batch_shape, y_dim)
g = matrix_and_mvn_to_gamma_gaussian(matrix, y_mvn)
xy = torch.randn(sample_shape + batch_shape + (x_dim + y_dim, ))
s = torch.rand(sample_shape + batch_shape)
actual_log_prob = g.log_density(xy, s)
x, y = xy[..., :x_dim], xy[..., x_dim:]
y_pred = x.unsqueeze(-2).matmul(matrix).squeeze(-2)
loc = y_pred + y_mvn.loc
scaled_prec = y_mvn.precision_matrix * s.unsqueeze(-1).unsqueeze(-1)
expected_log_prob = dist.MultivariateNormal(
loc, precision_matrix=scaled_prec).log_prob(y)
assert_close(actual_log_prob, expected_log_prob)