def test_beta_binomial_dependent_sample(): total = 10 counts = dist.Binomial(total, 0.3).sample() concentration1 = torch.tensor(0.5) concentration0 = torch.tensor(1.5) prior = dist.Beta(concentration1, concentration0) posterior = dist.Beta(concentration1 + counts, concentration0 + total - counts) def model(counts): prob = pyro.sample("prob", prior) pyro.sample("counts", dist.Binomial(total, prob), obs=counts) reparam_model = poutine.reparam( model, { "prob": ConjugateReparam( lambda counts: dist.Beta(1 + counts, 1 + total - counts)), }, ) with poutine.trace() as tr, pyro.plate("particles", 10000): reparam_model(counts) samples = tr.trace.nodes["prob"]["value"] assert_close(samples.mean(), posterior.mean, atol=0.01) assert_close(samples.std(), posterior.variance.sqrt(), atol=0.01)
def test_beta_binomial_hmc(): num_samples = 1000 total = 10 counts = dist.Binomial(total, 0.3).sample() concentration1 = torch.tensor(0.5) concentration0 = torch.tensor(1.5) prior = dist.Beta(concentration1, concentration0) likelihood = dist.Beta(1 + counts, 1 + total - counts) posterior = dist.Beta(concentration1 + counts, concentration0 + total - counts) def model(): prob = pyro.sample("prob", prior) pyro.sample("counts", dist.Binomial(total, prob), obs=counts) reparam_model = poutine.reparam(model, {"prob": ConjugateReparam(likelihood)}) kernel = HMC(reparam_model) samples = MCMC(kernel, num_samples, warmup_steps=0).run() pred = Predictive(reparam_model, samples, num_samples=num_samples) trace = pred.get_vectorized_trace() samples = trace.nodes["prob"]["value"] assert_close(samples.mean(), posterior.mean, atol=0.01) assert_close(samples.std(), posterior.variance.sqrt(), atol=0.01)
def guide(prior, obs, num_obs): a = pyro.param("a", prior["A"], constraint=constraints.positive) pyro.sample("p_A", dist.Beta(a[0], a[1])) b = pyro.param("b", prior["B"], constraint=constraints.positive) pyro.sample("p_B", dist.Beta(b[:, 0], b[:, 1]).to_event(1)) c = pyro.param("c", prior["C"], constraint=constraints.positive) pyro.sample("p_C", dist.Beta(c[:, 0], c[:, 1]).to_event(1))
def _update_the_policy_network(self, state_batch_tensor, advantages, action_batch_tensor): action_alpha_old, action_beta_old = self.policy_network( state_batch_tensor) old_beta_dist = dist.Beta(action_alpha_old, action_beta_old) old_action_prob = old_beta_dist.batch_log_pdf(action_batch_tensor) old_action_prob = old_action_prob.detach() if self.use_cuda: advantages = Variable(advantages).cuda() else: advantages = Variable(advantages) for _ in range(self.args.policy_update_step): action_alpha, action_beta = self.policy_network(state_batch_tensor) new_beta_dist = dist.Beta(action_alpha, action_beta) new_action_prob = new_beta_dist.batch_log_pdf(action_batch_tensor) # calculate the ratio ratio = torch.exp(new_action_prob - old_action_prob) # calculate the surr surr1 = ratio * advantages surr2 = torch.clamp(ratio, 1 - self.args.epsilon, 1 + self.args.epsilon) * advantages loss_policy = -torch.min(surr1, surr2).mean() self.optimizer_policy.zero_grad() loss_policy.backward() # update self.optimizer_policy.step() return loss_policy
def guide(prior, obs, num_obs): a = pyro.param('a', prior['A'], constraint=constraints.positive) pyro.sample('p_A', dist.Beta(a[0], a[1])) b = pyro.param('b', prior['B'], constraint=constraints.positive) pyro.sample('p_B', dist.Beta(b[:, 0], b[:, 1]).to_event(1)) c = pyro.param('c', prior['C'], constraint=constraints.positive) pyro.sample('p_C', dist.Beta(c[:, 0], c[:, 1]).to_event(1))
def model(self, *args, **kwargs): I = self._data['segments'] pi = pyro.sample("pi", dist.Dirichlet(self._params['init_probs'])) probs_z = pyro.sample("cnv_probs", dist.Dirichlet((1- self._params['t']) * torch.eye(self._params['hidden_dim']) + ( self._params['t'])).to_event(1)) probs_y = torch.tensor([[2., 64., 32., 21.5, 16., 43.],[64., 64., 64., 64., 64., 64.]]) z = pyro.sample("z_0", dist.Categorical(pi), infer={"enumerate": "parallel"}) pyro.sample("y_{}".format(0), dist.Beta(probs_y[0, z], probs_y[1, z]), obs=self._data['data'][0, 0]) for i in pyro.markov(range(1,I)): z = pyro.sample("z_{}".format(i), dist.Categorical(Vindex(probs_z)[z]), infer={"enumerate": "parallel"}) pyro.sample("y_{}".format(i), dist.Beta(probs_y[0,z], probs_y[1,z]), obs= self._data['data'][i,0])
def model(y,x,N): w = pyro.sample('w'.format(''), dist.Beta(Variable(26.072914040168385*torch.ones([amb(1)])),Variable((42.3120851154)*torch.ones([amb(1)])))) with pyro.iarange('b_range_'.format(''), N): b = pyro.sample('b'.format(''), dist.Gamma(Variable((5.63887222899)*torch.ones([amb(N)])),Variable((40.1978121928)*torch.ones([amb(N)])))) with pyro.iarange('p_range_'.format(''), N): p = pyro.sample('p'.format(''), dist.Beta(Variable((52.1419233118)*torch.ones([amb(N)])),Variable((83.6618285099)*torch.ones([amb(N)])))) pyro.sample('obs__100'.format(), dist.Beta(w*x+b,p), obs=y)
def __init__(self, ode_op, ode_model): super(LNAGenModel, self).__init__() self._ode_op = ode_op self._ode_model = ode_model self.ode_params1 = PyroSample(dist.Beta(2, 1)) self.ode_params2 = PyroSample(dist.HalfNormal(1)) self.ode_params3 = PyroSample(dist.Beta(1, 2))
def _model_(self): control_prior = pyro.sample('control_p', dist.Beta(1, 1)) treatment_prior = pyro.sample('treatment_p', dist.Beta(1, 1)) return pyro.sample( 'obs', dist.Binomial(self.traffic_size, torch.stack([control_prior, treatment_prior])), obs=self.outcome)
def model(data, params): # initialize data data = {k: torch.tensor(v).float() for k, v in data.items()} n = data["n"] k = data["k"] # model block theta = pyro.sample("theta", dist.Beta(1., 1.)) thetaprior = pyro.sample("thetaprior", dist.Beta(1., 1.)) k = pyro.sample("k", dist.Binomial(n, theta), obs=k)
def __init__(self, ode_op, ode_model): super(PlantModel, self).__init__() self._ode_op = ode_op self._ode_model = ode_model # TODO: Incorporate appropriate priors (cf. MATALB codes from Daewook) self.ode_params1 = PyroSample(dist.Gamma(1, 1000)) # dG self.ode_params2 = PyroSample(dist.Gamma(1, 1000)) # dP self.ode_params3 = PyroSample(dist.Beta(0.5, 0.5)) # G0 self.ode_params4 = PyroSample(dist.Beta(0.5, 0.5)) # P0
def model(prior, obs, num_obs): p_A = pyro.sample('p_A', dist.Beta(1, 1)) p_B = pyro.sample('p_B', dist.Beta(torch.ones(2), torch.ones(2)).to_event(1)) p_C = pyro.sample('p_C', dist.Beta(torch.ones(2), torch.ones(2)).to_event(1)) with pyro.plate('data_plate', num_obs): A = pyro.sample('A', dist.Bernoulli(p_A.expand(num_obs)), obs=obs['A']) # Vindex used to ensure proper indexing into the enumerated sample sites B = pyro.sample('B', dist.Bernoulli(Vindex(p_B)[A.type(torch.long)]), infer={"enumerate": "parallel"}) pyro.sample('C', dist.Bernoulli(Vindex(p_C)[B.type(torch.long)]), obs=obs['C'])
def kl(self, dist_a, prior=None): if prior == None: # use standard reparamterizer return self._kld_beta_kerman_prior(dist_a['beta']['conc1'], dist_a['beta']['conc2']) # we have two distributions provided (eg: VRNN) return torch.sum( D.kl_divergence( PD.Beta(dist_a['beta']['conc1'], dist_a['beta']['conc2']), PD.Beta(prior['beta']['conc1'], prior['beta']['conc2'])), -1)
def model(data): # define the hyperparameters that control the beta prior alpha0 = torch.tensor(10.0) beta0 = torch.tensor(10.0) # sample f from the beta prior f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0)) pyro.sample("some", dist.Beta(torch.tensor(10.0), torch.tensor(5.0))) # loop over the observed data for i in range(len(data)): # observe datapoint i using the bernoulli likelihood pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])
def model(data, params): # initialize data data = {k: torch.tensor(v).float() for k, v in data.items()} n1 = data["n1"] n2 = data["n2"] k1 = data["k1"] k2 = data["k2"] # init parameters theta1 = pyro.sample("theta1", dist.Beta(1., 1.)) theta2 = pyro.sample("theta2", dist.Beta(1., 1.)) pyro.sample("k1", dist.Binomial(n1, theta1), obs=k1) pyro.sample("k2", dist.Binomial(n2, theta2), obs=k2)
def _kld_beta_kerman_prior(self, conc1, conc2): """ Internal function to do a KL-div against the prior. :param conc1: concentration 1. :param conc2: concentration 2. :returns: batch_size tensor of kld against prior. :rtype: torch.Tensor """ prior = PD.Beta(zeros_like(conc1) + 1 / 3, zeros_like(conc2) + 1 / 3) beta = PD.Beta(conc1, conc2) return torch.sum(D.kl_divergence(beta, prior), -1)
def test_beta_binomial(sample_shape, batch_shape): concentration1 = torch.randn(batch_shape).exp() concentration0 = torch.randn(batch_shape).exp() total = 10 obs = dist.Binomial(total, 0.2).sample(sample_shape + batch_shape) f = dist.Beta(concentration1, concentration0) g = dist.Beta(1 + obs, 1 + total - obs) fg, log_normalizer = f.conjugate_update(g) x = fg.sample(sample_shape) assert_close(f.log_prob(x) + g.log_prob(x), fg.log_prob(x) + log_normalizer)
def model(): d = dist.Bernoulli(p) context1 = pyro.plate("outer", 3, dim=-1) context2 = pyro.plate("inner", 2, dim=-2) pyro.sample("w", d) pyro.sample("b", dist.Beta(1.1, 1.1)) with context1: pyro.sample("x", d) with context2: pyro.sample("c", dist.Beta(1.1, 1.1)) pyro.sample("y", d) with context1, context2: pyro.sample("z", d)
def model_multi_obs_grp(obsmat): # some parameters can be directly derived from the data passed # K = 2 nparticipants = data.shape[0] nfeatures = data.shape[1] # number of rows in each person's matrix ncol = data.shape[2] # Background probability of different groups if tm.stickbreak: # stick breaking process for assigning weights to groups with pyro.plate("beta_plate", K - 1): beta_mix = pyro.sample("weights", dist.Beta(1, 10)) weights = tm.mix_weights(beta_mix) else: weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(tm.K))) # declare model parameters based on whether the data are row-normalized if tm.dtype == 'norm': pass # with pyro.plate('components', K): # # concentration parameters # concentration = pyro.sample('concentration', # dist.Gamma(2 * torch.ones(nfeatures,ncol), 1/3 * torch.ones(nfeatures,ncol)).to_event(2)) # # implementation for the dirichlet based model is not complete!!!! # with pyro.plat('data',obsmat.shape[0]): # assignment = pyro.sample('assignment', dist.Categorical(weights)) # #d = dist.Dirichlet(concentration[assignment,:,:].clone().detach()) # .detach() might interfere with backprop # d = dist.Dirichlet(concentration[assignment,i,:]) # pyro.sample('obs', d.to_event(1), obs=obsmat) elif tm.dtype == 'raw': with pyro.plate('components', tm.K): alphas = pyro.sample( 'alpha', dist.Gamma(2 * torch.ones(nfeatures, ncol), 1 / 3 * torch.ones(nfeatures, ncol)).to_event(2)) betas = pyro.sample( 'beta', dist.Gamma(2 * torch.ones(nfeatures, ncol), 1 / 3 * torch.ones(nfeatures, ncol)).to_event(2)) assignment = pyro.sample('assignment', dist.Categorical(weights)) # expand assignment to make dimensions match for r in np.arange(obsmat.shape[0]): rowind = obsmat[r, 1].type(torch.long) colind = obsmat[r, 2].type(torch.long) d = dist.Beta(alphas[assignment, rowind, colind], betas[assignment, rowind, colind]) pyro.sample('obs_{}'.format(r), d, obs=obsmat[r, 0])
def guide(y,x,N): arg_1 = torch.nn.Softplus()(pyro.param('arg_1', Variable(torch.ones((amb(1))), requires_grad=True))) arg_2 = torch.nn.Softplus()(pyro.param('arg_2', Variable(torch.ones((amb(1))), requires_grad=True))) w = pyro.sample('w'.format(''), dist.Beta(arg_1,arg_2)) arg_3 = torch.nn.Softplus()(pyro.param('arg_3', Variable(torch.ones((amb(N))), requires_grad=True))) arg_4 = torch.nn.Softplus()(pyro.param('arg_4', Variable(torch.ones((amb(N))), requires_grad=True))) with pyro.iarange('b_prange'): b = pyro.sample('b'.format(''), dist.Gamma(arg_3,arg_4)) arg_5 = torch.nn.Softplus()(pyro.param('arg_5', Variable(torch.ones((amb(N))), requires_grad=True))) arg_6 = torch.nn.Softplus()(pyro.param('arg_6', Variable(torch.ones((amb(N))), requires_grad=True))) with pyro.iarange('p_prange'): p = pyro.sample('p'.format(''), dist.Beta(arg_5,arg_6)) pass
def setUp(self): self.alpha = Variable(torch.Tensor([2.4])) self.beta = Variable(torch.Tensor([3.7])) self.test_data = Variable(torch.Tensor([0.4])) self.dist = dist.Beta(self.alpha, self.beta) self.batch_alpha = Variable(torch.Tensor([[2.4], [3.6]])) self.batch_beta = Variable(torch.Tensor([[3.7], [2.5]])) self.batch_test_data = Variable(torch.Tensor([[0.4], [0.6]])) self.batch_dist = dist.Beta(self.batch_alpha, self.batch_beta) self.analytic_mean = (self.alpha / (self.alpha + self.beta)) one = Variable(torch.ones([1])) self.analytic_var = torch.pow(self.analytic_mean, 2.0) * self.beta / ( self.alpha * (self.alpha + self.beta + one)) self.n_samples = 50000
def mutual_info(self, params, eps=1e-9): """ I(z_d; x) ~ H(z_prior, z_d) + H(z_prior) :param params: parameters of distribution :param eps: tolerance :returns: batch_size mutual information (prop-to) tensor. :rtype: torch.Tensor """ z_true = PD.Beta(params['beta']['conc1'], params['beta']['conc2']) z_match = PD.Beta(params['q_z_given_xhat']['beta']['conc1'], params['q_z_given_xhat']['beta']['conc2']) kl_proxy_to_xent = torch.sum(D.kl_divergence(z_match, z_true), dim=-1) return self.config['continuous_mut_info'] * kl_proxy_to_xent
def test_init(): total = 10 counts = dist.Binomial(total, 0.3).sample() concentration1 = torch.tensor(0.5) concentration0 = torch.tensor(1.5) prior = dist.Beta(concentration1, concentration0) likelihood = dist.Beta(1 + counts, 1 + total - counts) def model(): x = pyro.sample("x", prior) pyro.sample("counts", dist.Binomial(total, x), obs=counts) return x check_init_reparam(model, ConjugateReparam(likelihood))
def model_0(sequences, lengths, args, batch_size=None, include_prior=True): assert not torch._C._get_tracing_state() num_sequences, max_length, data_dim = sequences.shape with poutine.mask(mask=include_prior): # Our prior on transition probabilities will be: # stay in the same state with 90% probability; uniformly jump to another # state with 10% probability. probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1)) # We put a weak prior on the conditional probability of a tone sounding. # We know that on average about 4 of 88 tones are active, so we'll set a # rough weak prior of 10% of the notes being active at any one time. probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, data_dim]).to_event(2)) # In this first model we'll sequentially iterate over sequences in a # minibatch; this will make it easy to reason about tensor shapes. tones_plate = pyro.plate("tones", data_dim, dim=-1) for i in pyro.plate("sequences", len(sequences), batch_size): length = lengths[i] sequence = sequences[i, :length] x = 0 for t in pyro.markov(range(length)): # On the next line, we'll overwrite the value of x with an updated # value. If we wanted to record all x values, we could instead # write x[t] = pyro.sample(...x[t-1]...). x = pyro.sample("x_{}_{}".format(i, t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) with tones_plate: pyro.sample("y_{}_{}".format(i, t), dist.Bernoulli(probs_y[x.squeeze(-1)]), obs=sequence[t])
def model_3(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length hidden_dim = int(args.hidden_dim**0.5) # split between w and x with poutine.mask(mask=include_prior): probs_w = pyro.sample( "probs_w", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1)) probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(hidden_dim) + 0.1).to_event(1)) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([hidden_dim, hidden_dim, data_dim]).to_event(3)) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] w, x = 0, 0 for t in pyro.markov(range(max_length if args.jit else lengths.max())): with poutine.mask(mask=(t < lengths).unsqueeze(-1)): w = pyro.sample("w_{}".format(t), dist.Categorical(probs_w[w]), infer={"enumerate": "parallel"}) x = pyro.sample("x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) with tones_plate as tones: pyro.sample("y_{}".format(t), dist.Bernoulli(probs_y[w, x, tones]), obs=sequences[batch, t])
def model_2(sequences, lengths, args, batch_size=None, include_prior=True): with ignore_jit_warnings(): num_sequences, max_length, data_dim = map(int, sequences.shape) assert lengths.shape == (num_sequences, ) assert lengths.max() <= max_length with poutine.mask(mask=include_prior): probs_x = pyro.sample( "probs_x", dist.Dirichlet(0.9 * torch.eye(args.hidden_dim) + 0.1).to_event(1)) probs_y = pyro.sample( "probs_y", dist.Beta(0.1, 0.9).expand([args.hidden_dim, 2, data_dim]).to_event(3)) tones_plate = pyro.plate("tones", data_dim, dim=-1) with pyro.plate("sequences", num_sequences, batch_size, dim=-2) as batch: lengths = lengths[batch] x, y = 0, 0 for t in pyro.markov(range(max_length if args.jit else lengths.max())): with poutine.mask(mask=(t < lengths).unsqueeze(-1)): x = pyro.sample("x_{}".format(t), dist.Categorical(probs_x[x]), infer={"enumerate": "parallel"}) # Note the broadcasting tricks here: to index probs_y on tensors x and y, # we also need a final tensor for the tones dimension. This is conveniently # provided by the plate associated with that dimension. with tones_plate as tones: y = pyro.sample("y_{}".format(t), dist.Bernoulli(probs_y[x, y, tones]), obs=sequences[batch, t]).long()
def __init__(self, ode_op, ode_model): super(ProteinGenModel, self).__init__() self._ode_op = ode_op self._ode_model = ode_model self.ode_params = PyroSample( dist.Beta(torch.tensor([1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), 2.0).to_event(1))
def _beta_bernoulli(data): alpha = torch.tensor([1.1, 1.1]) beta = torch.tensor([1.1, 1.1]) p_latent = pyro.sample('p_latent', dist.Beta(alpha, beta)) with pyro.plate('data', data.shape[0], dim=-2): pyro.sample('obs', dist.Bernoulli(p_latent), obs=data) return p_latent
def negative_binomial_guide( num_sites, num_days, num_predictors, predictors, data=None ): # Parameters for p alpha_0 = pyro.param( "alpha_0", 5 * torch.ones(num_sites), constraint=positive ) alpha_1 = pyro.param( "alpha_1", 5 * torch.ones(num_sites), constraint=positive ) means_epsilon = pyro.param("means_epsilon", torch.zeros(num_sites)) means_betas = pyro.param("means_betas", torch.zeros(num_predictors)) variance_epsilon = pyro.param( "variance_epsilon", torch.ones(num_sites), constraint=positive ) variance_betas = pyro.param( "variance_beta", torch.ones(num_predictors), constraint=positive ) with plate("beta_plates", num_predictors): pyro.sample("betas", dist.Normal(means_betas, variance_betas)) with plate("sites", size=num_sites): pyro.sample("epsilon", dist.Normal(means_epsilon, variance_epsilon)) pyro.sample("p", dist.Beta(alpha_0, alpha_1))
def model(data): y_prob = pyro.sample("y_prob", dist.Beta(1.0, 1.0)) y = pyro.sample("y", dist.Bernoulli(y_prob)) with pyro.plate("data", data.shape[0]): z = pyro.sample("z", dist.Bernoulli(0.65 * y + 0.1)) pyro.sample("obs", dist.Normal(2.0 * z, 1.0), obs=data) pyro.sample("nuisance", dist.Bernoulli(0.3))