def filter(self, value): """ Compute posterior over final state given a sequence of observations. :param ~torch.Tensor value: A sequence of observations. :return: A posterior distribution over latent states at the final time step, represented as a pair ``(cat, mvn)``, where :class:`~pyro.distributions.Categorical` distribution over mixture components and ``mvn`` is a :class:`~pyro.distributions.MultivariateNormal` with rightmost batch dimension ranging over mixture components. This can then be used to initialize a sequential Pyro model for prediction. :rtype: tuple """ ndims = max(len(self.batch_shape), value.dim() - 2) time = Variable("time", Bint[self.event_shape[0]]) value = tensor_to_funsor(value, ("time", ), 1) seq_sum_prod = naive_sequential_sum_product if self.exact else sequential_sum_product with interpretation(eager if self.exact else moment_matching): logp = self._trans + self._obs(value=value) logp = seq_sum_prod(ops.logaddexp, ops.add, logp, time, { "class": "class(time=1)", "state": "state(time=1)" }) logp += self._init logp = logp.reduce(ops.logaddexp, frozenset(["class", "state"])) cat, mvn = funsor_to_cat_and_mvn(logp, ndims, ("class(time=1)", )) cat = cat.expand(self.batch_shape) mvn = mvn.expand(self.batch_shape + cat.logits.shape[-1:]) return cat, mvn
def test_funsor_to_cat_and_mvn(batch_shape, event_shape, int_size, real_size): logits = torch.randn(batch_shape + event_shape + (int_size, )) expected_cat = dist.Categorical(logits=logits) expected_mvn = random_mvn(batch_shape + event_shape + (int_size, ), real_size) event_dims = tuple("abc"[:len(event_shape)]) + ("component", ) ndims = len(expected_cat.batch_shape) funsor_ = (tensor_to_funsor(logits, event_dims) + dist_to_funsor(expected_mvn, event_dims)(value="value")) assert isinstance(funsor_, Funsor) actual_cat, actual_mvn = funsor_to_cat_and_mvn(funsor_, ndims, event_dims) assert isinstance(actual_cat, dist.Categorical) assert isinstance(actual_mvn, dist.MultivariateNormal) assert actual_cat.batch_shape == expected_cat.batch_shape assert actual_mvn.batch_shape == expected_mvn.batch_shape assert_close(actual_cat.logits, expected_cat.logits, atol=1e-4, rtol=None) assert_close(actual_mvn.loc, expected_mvn.loc, atol=1e-4, rtol=None) assert_close(actual_mvn.precision_matrix, expected_mvn.precision_matrix, atol=1e-4, rtol=None)
def filter_and_predict(self, data, smoothing=False): trans_logits, trans_probs, trans_mvn, obs_mvn, x_trans_dist, y_dist = self.get_tensors_and_dists( ) log_prob = funsor.Number(0.) s_vars = { -1: funsor.Tensor(torch.tensor(0), dtype=self.num_components) } x_vars = {-1: None} predictive_x_dists, predictive_y_dists, filtering_dists = [], [], [] test_LLs = [] for t, y in enumerate(data): s_vars[t] = funsor.Variable(f's_{t}', funsor.bint(self.num_components)) x_vars[t] = funsor.Variable(f'x_{t}', funsor.reals(self.hidden_dim)) log_prob += dist.Categorical(trans_probs(s=s_vars[t - 1]), value=s_vars[t]) if t == 0: log_prob += self.x_init_mvn(value=x_vars[t]) else: log_prob += x_trans_dist(s=s_vars[t], x=x_vars[t - 1], y=x_vars[t]) if t > 0: log_prob = log_prob.reduce( ops.logaddexp, frozenset([s_vars[t - 1].name, x_vars[t - 1].name])) # do 1-step prediction and compute test LL if t > 0: predictive_x_dists.append(log_prob) _log_prob = log_prob - log_prob.reduce(ops.logaddexp) predictive_y_dist = y_dist(s=s_vars[t], x=x_vars[t]) + _log_prob test_LLs.append( predictive_y_dist(y=y).reduce(ops.logaddexp).data.item()) predictive_y_dist = predictive_y_dist.reduce( ops.logaddexp, frozenset([f"x_{t}", f"s_{t}"])) predictive_y_dists.append( funsor_to_mvn(predictive_y_dist, 0, ())) log_prob += y_dist(s=s_vars[t], x=x_vars[t], y=y) # save filtering dists for forward-backward smoothing if smoothing: filtering_dists.append(log_prob) # do the backward recursion using previously computed ingredients if smoothing: # seed the backward recursion with the filtering distribution at t=T smoothing_dists = [filtering_dists[-1]] T = data.size(0) s_vars = { t: funsor.Variable(f's_{t}', funsor.bint(self.num_components)) for t in range(T) } x_vars = { t: funsor.Variable(f'x_{t}', funsor.reals(self.hidden_dim)) for t in range(T) } # do the backward recursion. # let p[t|t-1] be the predictive distribution at time step t. # let p[t|t] be the filtering distribution at time step t. # let f[t] denote the prior (transition) density at time step t. # then the smoothing distribution p[t|T] at time step t is # given by the following recursion. # p[t-1|T] = p[t-1|t-1] <p[t|T] f[t] / p[t|t-1]> # where <...> denotes integration of the latent variables at time step t. for t in reversed(range(T - 1)): integral = smoothing_dists[-1] - predictive_x_dists[t] integral += dist.Categorical(trans_probs(s=s_vars[t]), value=s_vars[t + 1]) integral += x_trans_dist(s=s_vars[t], x=x_vars[t], y=x_vars[t + 1]) integral = integral.reduce( ops.logaddexp, frozenset([s_vars[t + 1].name, x_vars[t + 1].name])) smoothing_dists.append(filtering_dists[t] + integral) # compute predictive test MSE and predictive variances predictive_means = torch.stack([d.mean for d in predictive_y_dists ]) # T-1 ydim predictive_vars = torch.stack([ d.covariance_matrix.diagonal(dim1=-1, dim2=-2) for d in predictive_y_dists ]) predictive_mse = (predictive_means - data[1:, :]).pow(2.0).mean(-1) if smoothing: # compute smoothed mean function smoothing_dists = [ funsor_to_cat_and_mvn(d, 0, (f"s_{t}", )) for t, d in enumerate(reversed(smoothing_dists)) ] means = torch.stack([d[1].mean for d in smoothing_dists]) # T 2 xdim means = torch.matmul(means.unsqueeze(-2), self.observation_matrix).squeeze( -2) # T 2 ydim probs = torch.stack([d[0].logits for d in smoothing_dists]).exp() probs = probs / probs.sum(-1, keepdim=True) # T 2 smoothing_means = (probs.unsqueeze(-1) * means).sum(-2) # T ydim smoothing_probs = probs[:, 1] return predictive_mse, torch.tensor(np.array(test_LLs)), predictive_means, predictive_vars, \ smoothing_means, smoothing_probs else: return predictive_mse, torch.tensor(np.array(test_LLs))