def model(data): log_prob = funsor.Number(0.) # s is the discrete latent state, # x is the continuous latent state, # y is the observed state. s_curr = funsor.Tensor(torch.tensor(0), dtype=2) x_curr = funsor.Tensor(torch.tensor(0.)) for t, y in enumerate(data): s_prev = s_curr x_prev = x_curr # A delayed sample statement. s_curr = funsor.Variable('s_{}'.format(t), funsor.bint(2)) log_prob += dist.Categorical(trans_probs[s_prev], value=s_curr) # A delayed sample statement. x_curr = funsor.Variable('x_{}'.format(t), funsor.reals()) log_prob += dist.Normal(x_prev, trans_noise[s_curr], value=x_curr) # Marginalize out previous delayed sample statements. if t > 0: log_prob = log_prob.reduce(ops.logaddexp, {s_prev.name, x_prev.name}) # An observe statement. log_prob += dist.Normal(x_curr, emit_noise, value=y) log_prob = log_prob.reduce(ops.logaddexp) return log_prob
def model(data): log_prob = funsor.to_funsor(0.) xs_curr = [funsor.Tensor(torch.tensor(0.)) for var in var_names] for t, y in enumerate(data): xs_prev = xs_curr # A delayed sample statement. xs_curr = [ funsor.Variable(name + '_{}'.format(t), funsor.reals()) for name in var_names ] for i, x_curr in enumerate(xs_curr): log_prob += dist.Normal(trans_eqs[var_names[i]](xs_prev), torch.exp(trans_noises[i]), value=x_curr) if t > 0: log_prob = log_prob.reduce( ops.logaddexp, frozenset([x_prev.name for x_prev in xs_prev])) # An observe statement. log_prob += dist.Normal(emit_eq(xs_curr), torch.exp(emit_noise), value=y) # Marginalize out all remaining delayed variables. return log_prob.reduce(ops.logaddexp), log_prob.gaussian
def __call__(self): # calls pyro.param so that params are exposed and constraints applied # should not create any new torch.Tensors after __init__ self.initialize_params() N_c = self.config["sizes"]["group"] N_s = self.config["sizes"]["individual"] log_prob = Tensor(torch.tensor(0.), OrderedDict()) plate_g = Tensor(torch.zeros(N_c), OrderedDict([("g", bint(N_c))])) plate_i = Tensor(torch.zeros(N_s), OrderedDict([("i", bint(N_s))])) if self.config["group"]["random"] == "continuous": eps_g_dist = plate_g + dist.Normal(**self.params["eps_g"])( value="eps_g") log_prob += eps_g_dist # individual-level random effects if self.config["individual"]["random"] == "continuous": eps_i_dist = plate_g + plate_i + dist.Normal( **self.params["eps_i"])(value="eps_i") log_prob += eps_i_dist return log_prob
def test_normal_gaussian_3(batch_shape): batch_dims = ('i', 'j', 'k')[:len(batch_shape)] inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape)) loc = Tensor(torch.randn(batch_shape), inputs) scale = Tensor(torch.randn(batch_shape).exp(), inputs) value = Tensor(torch.randn(batch_shape), inputs) expected = dist.Normal(loc, scale, value) assert isinstance(expected, Tensor) check_funsor(expected, inputs, reals()) g = dist.Normal(Variable('loc', reals()), scale, 'value') assert isinstance(g, Contraction) actual = g(loc=loc, value=value) check_funsor(actual, inputs, reals()) assert_close(actual, expected, atol=1e-4)
def test_normal_independent(): loc = random_tensor(OrderedDict(), reals(2)) scale = random_tensor(OrderedDict(), reals(2)).exp() fn = dist.Normal(loc['i'], scale['i'], value='z_i') assert fn.inputs['z_i'] == reals() d = Independent(fn, 'z', 'i', 'z_i') assert d.inputs['z'] == reals(2) sample = d.sample(frozenset(['z'])) assert isinstance(sample, Contraction) assert sample.inputs['z'] == reals(2)
def model(data): log_prob = funsor.to_funsor(0.) x_curr = funsor.Tensor(torch.tensor(0.)) for t, y in enumerate(data): x_prev = x_curr # A delayed sample statement. x_curr = funsor.Variable('x_{}'.format(t), funsor.reals()) log_prob += dist.Normal(1 + x_prev / 2., trans_noise, value=x_curr) # Optionally marginalize out the previous state. if t > 0 and not args.lazy: log_prob = log_prob.reduce(ops.logaddexp, x_prev.name) # An observe statement. log_prob += dist.Normal(0.5 + 3 * x_curr, emit_noise, value=y) # Marginalize out all remaining delayed variables. log_prob = log_prob.reduce(ops.logaddexp) return log_prob
def test_normal_affine(expr): scale = Tensor(torch.tensor(0.3), OrderedDict()) x = Variable('x', reals()) y = Variable('y', reals()) expected = dist.Normal(x, scale, y) actual = eval(expr) assert isinstance(actual, Contraction) assert dict(actual.inputs) == dict(expected.inputs), (actual.inputs, expected.inputs) for ta, te in zip(actual.terms, expected.terms): assert_close(ta.align(tuple(te.inputs)), te)
def loss_function(data, subsample_scale): # Lazily sample from the guide. loc, scale = encode(data) q = funsor.Independent(dist.Normal(loc['i'], scale['i'], value='z_i'), 'z', 'i', 'z_i') # Evaluate the model likelihood at the lazy value z. probs = decode('z') p = dist.Bernoulli(probs['x', 'y'], value=data['x', 'y']) p = p.reduce(ops.add, {'x', 'y'}) # Construct an elbo. This is where sampling happens. elbo = funsor.Integrate(q, p - q, frozenset(['z'])) elbo = elbo.reduce(ops.add, 'batch') * subsample_scale loss = -elbo return loss
def update(p_x_tp1, t, y, var_names, emit_eq, emit_noise): """Computes p(x_{t+1} | y_{t+1}) from p(x_{t+1}). This is useful for iterating 1-step ahead predictions""" log_prob = p_x_tp1 x_tp1s = [ funsor.Variable(name + '_{}'.format(t + 1), funsor.reals()) for name in var_names ] log_p_x = log_prob log_prob += dist.Normal(emit_eq(x_tp1s), emit_noise, value=y) log_p_y = log_prob.reduce(ops.logaddexp, frozenset([x_tp1.name for x_tp1 in x_tp1s])) log_p_x_y = log_prob + log_p_x - log_p_y return log_p_x_y
def one_step_prediction(p_x_tp1, t, var_names, emit_eq, emit_noise): """Computes p(y_{t+1}) from p(x_{t+1}). We assume y_t is scalar, so only one emit_eq""" log_prob = p_x_tp1 x_tp1s = [ funsor.Variable(name + '_{}'.format(t + 1), funsor.reals()) for name in var_names ] y_tp1 = funsor.Variable('y_{}'.format(t + 1), funsor.reals()) log_prob += dist.Normal(emit_eq(x_tp1s), torch.exp(emit_noise), value=y_tp1) log_prob = log_prob.reduce(ops.logaddexp, frozenset([x_tp1.name for x_tp1 in x_tp1s])) return log_prob
def test_normal_affine(expr): scale = Tensor(torch.tensor(0.3), OrderedDict()) x = Variable('x', reals()) y = Variable('y', reals()) expected = dist.Normal(x, scale, y) actual = eval(expr) assert isinstance(actual, Joint) assert dict(actual.inputs) == dict(expected.inputs), (actual.inputs, expected.inputs) assert_close(actual.gaussian.align(tuple(expected.gaussian.inputs)), expected.gaussian) assert_close(actual.discrete.align(tuple(expected.discrete.inputs)), expected.discrete)
def next_state(p_x_t, t, var_names, trans_eqs, trans_noises): """Computes p(x_{t+1}) from p(x_t)""" log_prob = p_x_t x_ts = [ funsor.Variable(name + '_{}'.format(t), funsor.reals()) for name in var_names ] x_tp1s = [ funsor.Variable(name + '_{}'.format(t + 1), funsor.reals()) for name in var_names ] for i, x_tp1 in enumerate(x_tp1s): log_prob += dist.Normal(trans_eqs[var_names[i]](x_ts), torch.exp(trans_noises[i]), value=x_tp1) log_prob = log_prob.reduce(ops.logaddexp, frozenset([x_t.name for x_t in x_ts])) return log_prob
def test_normal_density(batch_shape): batch_dims = ('i', 'j', 'k')[:len(batch_shape)] inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape)) @funsor.of_shape(reals(), reals(), reals()) def normal(loc, scale, value): return -((value - loc)**2) / (2 * scale**2) - scale.log() - math.log( math.sqrt(2 * math.pi)) check_funsor(normal, { 'loc': reals(), 'scale': reals(), 'value': reals() }, reals()) loc = Tensor(torch.randn(batch_shape), inputs) scale = Tensor(torch.randn(batch_shape).exp(), inputs) value = Tensor(torch.randn(batch_shape), inputs) expected = normal(loc, scale, value) check_funsor(expected, inputs, reals()) actual = dist.Normal(loc, scale, value) check_funsor(actual, inputs, reals()) assert_close(actual, expected)
def test_normal_defaults(): loc = Variable('loc', reals()) scale = Variable('scale', reals()) value = Variable('value', reals()) assert dist.Normal(loc, scale) is dist.Normal(loc, scale, value)
def __call__(self): # calls pyro.param so that params are exposed and constraints applied # should not create any new torch.Tensors after __init__ self.initialize_params() N_state = self.config["sizes"]["state"] # initialize gamma to uniform gamma = Tensor( torch.zeros((N_state, N_state)), OrderedDict([("y_prev", bint(N_state))]), ) N_v = self.config["sizes"]["random"] N_c = self.config["sizes"]["group"] log_prob = [] plate_g = Tensor(torch.zeros(N_c), OrderedDict([("g", bint(N_c))])) # group-level random effects if self.config["group"]["random"] == "discrete": # group-level discrete effect e_g = Variable("e_g", bint(N_v)) e_g_dist = plate_g + dist.Categorical(**self.params["e_g"])(value=e_g) log_prob.append(e_g_dist) eps_g = (plate_g + self.params["eps_g"]["theta"])(e_g=e_g) elif self.config["group"]["random"] == "continuous": eps_g = Variable("eps_g", reals(N_state)) eps_g_dist = plate_g + dist.Normal(**self.params["eps_g"])(value=eps_g) log_prob.append(eps_g_dist) else: eps_g = to_funsor(0.) N_s = self.config["sizes"]["individual"] plate_i = Tensor(torch.zeros(N_s), OrderedDict([("i", bint(N_s))])) # individual-level random effects if self.config["individual"]["random"] == "discrete": # individual-level discrete effect e_i = Variable("e_i", bint(N_v)) e_i_dist = plate_g + plate_i + dist.Categorical( **self.params["e_i"] )(value=e_i) * self.raggedness_masks["individual"](t=0) log_prob.append(e_i_dist) eps_i = (plate_i + plate_g + self.params["eps_i"]["theta"](e_i=e_i)) elif self.config["individual"]["random"] == "continuous": eps_i = Variable("eps_i", reals(N_state)) eps_i_dist = plate_g + plate_i + dist.Normal(**self.params["eps_i"])(value=eps_i) log_prob.append(eps_i_dist) else: eps_i = to_funsor(0.) # add group-level and individual-level random effects to gamma gamma = gamma + eps_g + eps_i N_state = self.config["sizes"]["state"] # we've accounted for all effects, now actually compute gamma_y gamma_y = gamma(y_prev="y(t=1)") y = Variable("y", bint(N_state)) y_dist = plate_g + plate_i + dist.Categorical( probs=gamma_y.exp() / gamma_y.exp().sum() )(value=y) # observation 1: step size step_dist = plate_g + plate_i + dist.Gamma( **{k: v(y_curr=y) for k, v in self.params["step"].items()} )(value=self.observations["step"]) # step size zero-inflation if self.config["zeroinflation"]: step_zi = dist.Categorical(probs=self.params["zi_step"]["zi_param"](y_curr=y))( value="zi_step") step_zi_dist = plate_g + plate_i + dist.Delta(self.config["MISSING"], 0.)( value=self.observations["step"]) step_dist = (step_zi + Stack("zi_step", (step_dist, step_zi_dist))).reduce(ops.logaddexp, "zi_step") # observation 2: step angle angle_dist = plate_g + plate_i + dist.VonMises( **{k: v(y_curr=y) for k, v in self.params["angle"].items()} )(value=self.observations["angle"]) # observation 3: dive activity omega_dist = plate_g + plate_i + dist.Beta( **{k: v(y_curr=y) for k, v in self.params["omega"].items()} )(value=self.observations["omega"]) # dive activity zero-inflation if self.config["zeroinflation"]: omega_zi = dist.Categorical(probs=self.params["zi_omega"]["zi_param"](y_curr=y))( value="zi_omega") omega_zi_dist = plate_g + plate_i + dist.Delta(self.config["MISSING"], 0.)( value=self.observations["omega"]) omega_dist = (omega_zi + Stack("zi_omega", (omega_dist, omega_zi_dist))).reduce(ops.logaddexp, "zi_omega") # finally, construct the term for parallel scan reduction hmm_factor = step_dist + angle_dist + omega_dist hmm_factor = hmm_factor * self.raggedness_masks["individual"] hmm_factor = hmm_factor * self.raggedness_masks["timestep"] # copy masking behavior of pyro.infer.TraceEnum_ELBO._compute_model_factors hmm_factor = hmm_factor + y_dist log_prob.insert(0, hmm_factor) return log_prob