def guide(home_id, away_id, score1_obs=None, score2_obs=None): mu_locs = pyro.param("mu_loc", torch.tensor(0.0).expand(4)) mu_scales = pyro.param( "mu_scale", torch.tensor(0.1).expand(4), constraint=constraints.positive, ) pyro.sample("alpha", dist.Normal(mu_locs[0], mu_scales[0])) pyro.sample("sd_att", dist.LogNormal(mu_locs[1], mu_scales[1])) pyro.sample("sd_def", dist.LogNormal(mu_locs[2], mu_scales[2])) pyro.sample("home", dist.Normal(mu_locs[3], mu_scales[3])) nt = len(np.unique(home_id)) mu_team_locs = pyro.param("mu_team_loc", torch.tensor(0.0).expand(2, nt)) mu_team_scales = pyro.param( "mu_team_scale", torch.tensor(0.1).expand(2, nt), constraint=constraints.positive, ) with pyro.plate("plate_teams", nt): pyro.sample("attack", dist.Normal(mu_team_locs[0], mu_team_scales[0])) pyro.sample("defend", dist.Normal(mu_team_locs[1], mu_team_scales[1]))
def guide_ret(*args, **kwargs): I, N = self._data['data'].shape seg_mean = torch.mean(self._data['data'] / self._data['pld'].reshape([I, 1]), axis=1) batch = N if self._params['batch_size'] else self._params['batch_size'] param_weights = pyro.param("param_mixture_weights", lambda: torch.ones(self._params['K']) / self._params['K'], constraint=constraints.simplex) cnv_mean = pyro.param("param_cnv_probs", lambda: self.create_gaussian_init_values(), constraint=constraints.positive) cnv_var = pyro.param("param_cnv_var", lambda: torch.ones(I) * self._params['cnv_sd'], constraint=constraints.positive) seg_var = pyro.param("param_seg_var", lambda: torch.ones(I) * self._params['cnv_sd'], constraint=constraints.positive) seg_mean = pyro.param("param_seg_mean", lambda: seg_mean) gamma_scale = pyro.param("param_norm_factor", lambda: torch.sum(self._data['data'], axis = 0) / torch.sum(seg_mean), constraint=constraints.positive) pyro.sample('mixture_weights', dist.Dirichlet(param_weights)) with pyro.plate('segments', I): pyro.sample('segment_mean', dist.LogNormal(torch.log(seg_mean) - seg_var ** 2 / 2, seg_var)) with pyro.plate('components', self._params['K']): pyro.sample('cnv_probs', dist.LogNormal(torch.log(cnv_mean) - cnv_var ** 2 / 2, cnv_var)) with pyro.plate("data2", N, batch): pyro.sample('norm_factor', dist.Delta(gamma_scale))
def model(self, zero_data, covariates): data_dim = zero_data.size(-1) feature_dim = covariates.size(-1) bias = pyro.sample("bias", dist.Normal(5.0, 10.0).expand((data_dim,)).to_event(1)) weight = pyro.sample("weight", dist.Normal(0.0, 5).expand((feature_dim,)).to_event(1)) # We'll sample a time-global scale parameter outside the time plate, # then time-local iid noise inside the time plate. drift_scale = pyro.sample("drift_scale", dist.LogNormal(-20, 5.0).expand((1,)).to_event(1)) with self.time_plate: # We'll use a reparameterizer to improve variational fit. The model would still be # correct if you removed this context manager, but the fit appears to be worse. with poutine.reparam(config={"drift": LocScaleReparam()}): drift = pyro.sample("drift", dist.Normal(zero_data.double(), drift_scale.double()).to_event(1)) # After we sample the iid "drift" noise we can combine it in any time-dependent way. # It is important to keep everything inside the plate independent and apply dependent # transforms outside the plate. motion = drift.cumsum(-2) # A Brownian motion. # The prediction now includes three terms. prediction = motion + bias + (weight * covariates).sum(-1, keepdim=True) assert prediction.shape[-2:] == zero_data.shape # Construct the noise distribution and predict. noise_scale = pyro.sample("noise_scale", dist.LogNormal(0.0, 1.0).expand((1,)).to_event(1)) noise_dist = dist.Normal(0, noise_scale) self.predict(noise_dist, prediction)
def state_space_model(data, N=1, T=2, prior_drift=0., verbose=False): # global rvs drift = pyro.sample('drift', dist.Normal(prior_drift, 1)) vol = pyro.sample('vol', dist.LogNormal(0, 1)) uncert = pyro.sample('uncert', dist.LogNormal(-5, 1)) if verbose: print(f"Using drift = {drift}, vol = {vol}, uncert = {uncert}") # the latent time series you want to infer # since you want to output this, we initialize a vector where you'll # save the inferred values latent = torch.empty((T, N)) # +1 comes from hidden initial condition # I think you want to plate out the same state space model for N different obs with pyro.plate('data_plate', N) as n: x0 = pyro.sample('x0', dist.Normal(drift, vol)) # or whatever your IC might be latent[0, n] = x0 # now comes the markov part, as you correctly noted for t in pyro.markov(range(1, T)): x_t = pyro.sample(f"x_{t}", dist.Normal(latent[t - 1, n] + drift, vol)) y_t = pyro.sample(f"y_{t}", dist.Normal(x_t, uncert), obs=data[t - 1, n] if data is not None else None) latent[t, n] = x_t return pyro.deterministic('latent', latent)
def guide(self, X): _id = self._id K = self.K N, D = X.shape loc_locinit = self.param_init[f'loc_loc_init_{_id}'] loc_scaleinit = self.param_init[f'loc_scale_init_{_id}'] cov_diag_loc_init = self.param_init[f'cov_diag_loc_init_{_id}'] cov_diag_scale_init = self.param_init[f'cov_diag_scale_init_{_id}'] cov_factor_loc = self.param_init[f'cov_factor_loc_init_{_id}'] cov_factor_scale = self.param_init[f'cov_factor_scale_init_{_id}'] with pyro.plate(f'D_{_id}', D, dim=-1): loc_loc = pyro.param(f'loc_loc_{_id}', loc_locinit) loc_scale = pyro.param(f'loc_scale_{_id}', loc_scaleinit, constraint=constraints.positive) loc = pyro.sample(f'loc_{_id}', dist.LogNormal(loc_loc, loc_scale)) cov_diag_loc = pyro.param(f'cov_diag_loc_{_id}', cov_diag_loc_init) cov_diag_scale = pyro.param(f'cov_diag_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive) cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale)) cov_diag = cov_diag*torch.ones(D) cov_diag = cov_diag + jitter # sample variables with pyro.plate(f'K_{_id}', K): cov_factor_loc = pyro.param(f'cov_factor_loc_{_id}', cov_factor_loc_init) cov_factor_scale = pyro.param(f'cov_factor_scale_{_id}', cov_factor_scale_init, constraint=constraints.positive) cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc, cov_factor_scale)) cov_factor = cov_factor.transpose(-2,-1) return loc, cov_factor, cov_diag
def guide(self, X): _id = self._id N, D = X.shape global_shrinkage_loc_init = self.param_init[f'global_shrinkage_loc_init_{_id}'] global_shrinkage_scale_init = self.param_init[f'global_shrinkage_scale_init_{_id}'] local_shrinkage_loc_init = self.param_init[f'local_shrinkage_loc_init_{_id}'] local_shrinkage_scale_init = self.param_init[f'local_shrinkage_scale_init_{_id}'] cov_diag_loc_init = self.param_init[f'cov_diag_loc_init_{_id}'] cov_diag_scale_init = self.param_init[f'cov_diag_scale_init_{_id}'] cov_factor_loc_init = self.param_init[f'cov_factor_loc_init_{_id}'] global_shrinkage_loc = pyro.param(f'global_shrinkage_loc_{_id}', global_shrinkage_loc_init) global_shrinkage_scale = pyro.param(f'global_shrinkage_scale_{_id}', global_shrinkage_scale_init, constraint=constraints.positive) tau = pyro.sample(f'global_shrinkage_{_id}', dist.LogNormal(global_shrinkage_loc,global_shrinkage_scale)) b = pyro.sample('b', dist.InverseGamma(0.5,1./torch.ones(D)**2).to_event(1)) local_shrinkage_loc = pyro.param(f'local_shrinkage_loc_{_id}', local_shrinkage_loc_init) local_shrinkage_scale = pyro.param(f'local_shrinkage_scale_{_id}', local_shrinkage_scale_init, constraint=constraints.positive) lambdasquared = pyro.sample(f'local_shrinkage_{_id}', dist.LogNormal(local_shrinkage_loc,local_shrinkage_scale).to_event(1)) cov_diag_loc = pyro.param(f'cov_diag_loc_{_id}', cov_diag_loc_init) cov_diag_scale = pyro.param(f'cov_diag_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive) cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale).to_event(1)) cov_diag = cov_diag*torch.ones(D) lambdasquared = lambdasquared.squeeze() if lambdasquared.dim() == 1: cov_factor_scale = torch.ger(torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,))) else: cov_factor_scale = torch.einsum('bp, br->bpr', torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,))) cov_factor_loc = pyro.param(f'cov_factor_loc_{_id}', cov_factor_loc_init) cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc, cov_factor_scale).to_event(2)) return tau, lambdasquared, cov_factor, cov_diag
def model(self, X): _id = self._id K = self.K N, D = X.shape loc_locinit = self.param_init[f'loc_prior_loc_init_{_id}'] loc_scaleinit = self.param_init[f'loc_prior_scale_init_{_id}'] cov_diag_loc_init = self.param_init[f'cov_diag_prior_loc_init_{_id}'] cov_diag_scale_init = self.param_init[f'cov_diag_prior_scale_init_{_id}'] cov_factor_loc_init = self.param_init[f'cov_factor_prior_loc_init_{_id}'] cov_factor_scale_init = self.param_init[f'cov_factor_prior_scale_init_{_id}'] with pyro.plate(f'D_{_id}', D): loc_loc = pyro.param(f'loc_loc_prior_{_id}', loc_locinit) loc_scale = pyro.param(f'loc_scale_prior_{_id}', loc_scaleinit, constraint=constraints.positive) loc = pyro.sample(f'loc_{_id}', dist.LogNormal(loc_loc, loc_scale)) cov_diag_loc = pyro.param(f'cov_diag_prior_loc_{_id}', cov_diag_loc_init) cov_diag_scale = pyro.param(f'cov_diag_prior_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive) cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale)) cov_diag = cov_diag + jitter cov_factor = None with pyro.plate(f'K_{_id}', K): cov_factor_loc = pyro.param(f'cov_factor_prior_loc_{_id}', cov_factor_loc_init) cov_factor_scale = pyro.param(f'cov_factor_prior_scale_{_id}', cov_factor_scale_init, constraint=constraints.positive) cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc,cov_factor_scale)) cov_factor = cov_factor.transpose(-2,-1) with pyro.plate(f'N_{_id}', size=N, subsample_size=self.batch_size) as ind: X = pyro.sample('obs', dist.LowRankMultivariateNormal(loc, cov_factor=cov_factor, cov_diag=cov_diag), obs=X.index_select(0, ind)) return X
def model(self, zero_data, covariates): duration,data_dim = zero_data.shape feature_dim = covariates.size(-1) drift_stability = pyro.sample("drift_stability", dist.Uniform(1, 2)) drift_scale = pyro.sample("drift_scale", dist.LogNormal(-20, 5)) with pyro.plate("item", data_dim, dim=-2): bias = pyro.sample("bias", dist.Normal(5.0, 10.0)) with pyro.plate('features',feature_dim,dim=-1): weight = pyro.sample("weight", dist.Normal(0.0, 5)) # We'll sample a time-global scale parameter outside the time plate, # then time-local iid noise inside the time plate. with self.time_plate: # We combine two different reparameterizers: the inner SymmetricStableReparam # is needed for the Stable site, and the outer LocScaleReparam is optional but # appears to improve inference. with poutine.reparam(config={"drift": LocScaleReparam()}): with poutine.reparam(config={"drift": SymmetricStableReparam()}): drift = pyro.sample("drift", dist.Stable(drift_stability, 0, drift_scale)) motion = drift.cumsum(-1) # A Brownian motion. # The prediction now includes three terms. regression = torch.matmul(covariates, weight[...,None]) prediction = motion + bias + regression.sum(axis=-1) prediction = prediction.unsqueeze(-1).transpose(-1, -3) # Finally we can construct a noise distribution. # We will share parameters across all time series. obs_scale = pyro.sample("obs_scale", dist.LogNormal(-5, 5)) noise_dist = dist.Normal(loc=0.0, scale=obs_scale.unsqueeze(-1)) self.predict(noise_dist, prediction)
def model(self, zero_data, covariates): assert zero_data.size(-1) == 1 # univariate duration = zero_data.size(-2) time, feature = covariates[..., 0], covariates[..., 1:] bias = pyro.sample("bias", dist.Normal(0, 10)) # construct a linear trend; we know that the sales are increasing # through years, so a positive-support prior should be used here trend_coef = pyro.sample("trend", dist.LogNormal(-2, 1)) trend = trend_coef * time # set prior of weights of the remaining covariates weight = pyro.sample( "weight", dist.Normal(0, 1).expand([feature.size(-1)]).to_event(1)) regressor = (weight * feature).sum(-1) # encode the additive weekly seasonality with pyro.plate("day_of_week", 7, dim=-1): seasonal = pyro.sample("seasonal", dist.Normal(0, 5)) seasonal = periodic_repeat(seasonal, duration, dim=-1) # make prediction prediction = bias + trend + seasonal + regressor # because Pyro forecasting framework is multivariate, # for univariate timeseries we need to make sure that # the last dimension is 1 prediction = prediction.unsqueeze(-1) # Now, we will use heavy tail noise because the data has some outliers # (such as Christmas day) dof = pyro.sample("dof", dist.Uniform(1, 10)) noise_scale = pyro.sample("noise_scale", dist.LogNormal(-2, 1)) noise_dist = dist.StudentT(dof.unsqueeze(-1), 0, noise_scale.unsqueeze(-1)) self.predict(noise_dist, prediction)
def _param_map_estimates( self, data: torch.Tensor, chi_ambient: torch.Tensor) -> Dict[str, torch.Tensor]: """Calculate MAP estimates of mu, the mean of the true count matrix, and lambda, the rate parameter of the Poisson background counts. Args: data: Dense tensor minibatch of cell by gene count data. chi_ambient: Point estimate of inferred ambient gene expression. Returns: mu_map: Dense tensor of Negative Binomial means for true counts. lambda_map: Dense tensor of Poisson rate params for noise counts. alpha_map: Dense tensor of Dirichlet concentration params that inform the overdispersion of the Negative Binomial. """ # Encode latents. enc = self.vi_model.encoder.forward(x=data, chi_ambient=chi_ambient) z_map = enc['z']['loc'] chi_map = self.vi_model.decoder.forward(z_map) phi_loc = pyro.param('phi_loc') phi_scale = pyro.param('phi_scale') phi_conc = phi_loc.pow(2) / phi_scale.pow(2) phi_rate = phi_loc / phi_scale.pow(2) alpha_map = 1. / dist.Gamma(phi_conc, phi_rate).mean y = (enc['p_y'] > 0).float() d_empty = dist.LogNormal(loc=pyro.param('d_empty_loc'), scale=pyro.param('d_empty_scale')).mean d_cell = dist.LogNormal(loc=enc['d_loc'], scale=pyro.param('d_cell_scale')).mean epsilon = dist.Gamma(enc['epsilon'] * self.vi_model.epsilon_prior, self.vi_model.epsilon_prior).mean if self.vi_model.include_rho: rho = pyro.param("rho_alpha") / (pyro.param("rho_alpha") + pyro.param("rho_beta")) else: rho = None # Calculate MAP estimates of mu and lambda. mu_map = self.vi_model.calculate_mu(epsilon=epsilon, d_cell=d_cell, chi=chi_map, y=y, rho=rho) lambda_map = self.vi_model.calculate_lambda( epsilon=epsilon, chi_ambient=chi_ambient, d_empty=d_empty, y=y, d_cell=d_cell, rho=rho, chi_bar=self.vi_model.avg_gene_expression) return {'mu': mu_map, 'lam': lambda_map, 'alpha': alpha_map}
def run_bo_pyro(params_bo, params_data, outer_loop_steps=10, q_size=5): """ run the bo """ X = params_data['X'] y = params_data['y'] #import ipdb; ipdb.set_trace() gpmodel = gp.models.GPRegression( X, y, gp.kernels.Matern52(input_dim=params_data['dim']), noise=torch.tensor(0.1), jitter=10**(-1)) gpmodel.kernel.set_prior("lengthscale", dist.LogNormal(0.0, 1.0)) gpmodel.kernel.set_prior("variance", dist.LogNormal(0.0, 1.0)) sampling_type = params_bo['sampling_type'] sample_size = params_bo['sample_size'] f_target = params_data['f_target'] print('run model with %s and %s samples' % (sampling_type, sample_size)) start_time = time.time() optimizer = torch.optim.Adam(gpmodel.parameters(), lr=0.001) gp.util.train(gpmodel, optimizer) #gpmodel.optimize() y_list_min = [] x_list_min = [] for i in range(outer_loop_steps): print('approach %s, sample size %s, outer loop step %s of %s' % (sampling_type, sample_size, i, outer_loop_steps)) xmin = next_x(gpmodel, sampling_type=sampling_type, sample_size=sample_size, q_size=q_size) print("next points evaluated:") #import ipdb; ipdb.set_trace() print(xmin, f_target(xmin)) # update the posterior with the new point found #TODO: do we evaluate the target function at all points?? update_posterior(xmin, f_target, gpmodel) # add to the list of best sampled points val_y, ind = gpmodel.y.min(0) val_x = gpmodel.X[ind, :] y_list_min.append(val_y.detach().numpy()) x_list_min.append(val_x.detach().numpy()) print("best points so far:") print(x_list_min[-1], y_list_min[-1]) print("run time %s seconds" % (time.time() - start_time)) res_dict = { 'X_exp': gpmodel.X.detach().numpy(), 'y_exp': gpmodel.y.numpy(), 'X_min': np.array(x_list_min), 'y_min': np.array(y_list_min) } return gpmodel, res_dict
def module_ppca_gm_means_sigma(self, input_batch, epsilon): max_hidden_dim = max(self.hidden_dims) batch_size = input_batch.shape[0] ppca_means = input_batch.new_zeros( [self.K, batch_size, self.output_dim]) if self.likelihood == 'normal': #with pyro.plate('ppca_sigma_plate', self.K): if self.group_isotropic: ppca_gm_sigma = pyro.sample( f'ppca_gm_sigma', dist.LogNormal(0, input_batch.new_ones(1, 1)).independent(1)) else: ppca_gm_sigma_list = [] ppca_gm_sigma = input_batch.new_ones(1, 1) for i in range(self.d): ppca_gm_sigma_list.append( pyro.sample( f'ppca_gm_sigma_{i}', dist.LogNormal(0, input_batch.new_ones( 1, self.n[i])).independent(1))) ppca_gm_sigma = torch_utils.krp_cw_torch( ppca_gm_sigma_list[i], ppca_gm_sigma, column=False) #with pyro.plate('alpha_plate', self.K): alpha_gm = pyro.sample( f'alpha_gm', dist.LogNormal(0, input_batch.new_ones([1, self.group_hidden_dim ])).independent(1)) if self.group_iterm is None: ppca_gm_means = self.group_term.multi_project( input_batch, remove_bias=False, tensorize=False #fast_svd=False ) else: ppca_gm_means = self.group_term( self.group_iterm( torch_utils.flatten_torch(input_batch), T=True #, fast_svd=False )) ppca_gm_means += self.group_term(epsilon[:, :self.group_hidden_dim] * alpha_gm[:, :self.group_hidden_dim]) if self.likelihood == 'bernoulli': return ppca_gm_means.sigmoid() if self.likelihood == 'normal': return ppca_gm_means, ppca_gm_sigma raise ValueError
def normal_model(): zero = torch.zeros(2) a = pyro.sample("a", dist.Normal(0, 1)) b = pyro.sample("b", dist.LogNormal(0, 1)) c = pyro.sample("c", dist.Normal(a, b)) d = pyro.sample("d", dist.LogNormal(a, b)) e = pyro.sample("e", dist.Normal(zero, b).to_event(1)) f = pyro.sample("f", dist.LogNormal(zero, b).to_event(1)) g = pyro.sample("g", dist.Normal(0, 1), obs=a) h = pyro.sample("h", dist.LogNormal(0, 1), obs=b) with pyro.plate("plate", 5): i = pyro.sample("i", dist.Normal(a, b)) j = pyro.sample("j", dist.LogNormal(a, b)) return a, b, c, d, e, f, g, h, i, j
def get_model(self, train_data, **kwargs): X_train, y_train = train_data # initialize the kernel and model pyro.clear_param_store() kernel = self.kernel(input_dim=X_train.shape[1]) gpr = gp.models.GPRegression(X_train, y_train, kernel, noise=torch.tensor(1.).double()) # optional: fit the model using MAP gpr.kernel.lengthscale = pyro.nn.PyroSample(dist.LogNormal(0.0, 1.0)) gpr.kernel.variance = pyro.nn.PyroSample(dist.LogNormal(0.0, 1.0)) return gpr
def guide(self, x, y=None): pyro.module("scanvi", self) with pyro.plate("batch", len(x)), poutine.scale(scale=self.scale_factor): z2_loc, z2_scale, l_loc, l_scale = self.z2l_encoder(x) pyro.sample("l", dist.LogNormal(l_loc, l_scale).to_event(1)) z2 = pyro.sample("z2", dist.Normal(z2_loc, z2_scale).to_event(1)) y_logits = self.classifier(z2) y_dist = dist.OneHotCategorical(logits=y_logits) if y is None: # x is unlabeled so sample y using q(y|z2) y = pyro.sample("y", y_dist) else: # x is labeled so add a classification loss term # (this way q(y|z2) learns from both labeled and unlabeled data) classification_loss = y_dist.log_prob(y) # Note that the negative sign appears because we're adding this term in the guide # and the guide log_prob appears in the ELBO as -log q pyro.factor( "classification_loss", -self.alpha * classification_loss, has_rsample=False, ) z1_loc, z1_scale = self.z1_encoder(z2, y) pyro.sample("z1", dist.Normal(z1_loc, z1_scale).to_event(1))
def model(self, peak_idx, read_depth, onehot_obs=None): pyro.module("decoder", self.decoder) #pyro.module("decoder", self.decoder) with pyro.plate("cells", peak_idx.shape[0]): # Dirichlet prior 𝑝(𝜃|𝛼) is replaced by a log-normal distribution theta_loc = self.prior_mu * peak_idx.new_ones( (peak_idx.shape[0], self.num_topics)) theta_scale = self.prior_std * peak_idx.new_ones( (peak_idx.shape[0], self.num_topics)) theta = pyro.sample( "theta", dist.LogNormal(theta_loc, theta_scale).to_event(1)) theta = theta / theta.sum(-1, keepdim=True) # conditional distribution of 𝑤𝑛 is defined as # 𝑤𝑛|𝛽,𝜃 ~ Categorical(𝜎(𝛽𝜃)) peak_probs = self.decoder(theta) pyro.sample( 'obs', dist.Multinomial( total_count=read_depth if onehot_obs is None else 1, probs=peak_probs).to_event(1), obs=onehot_obs)
def stable_model(): zero = torch.zeros(2) a = pyro.sample("a", dist.Normal(0, 1)) b = pyro.sample("b", dist.LogNormal(0, 1)) c = pyro.sample("c", dist.Stable(1.5, 0.0, b, a)) d = pyro.sample("d", dist.Stable(1.5, 0.0, b, 0.0), obs=a) e = pyro.sample("e", dist.Stable(1.5, 0.1, b, a)) f = pyro.sample("f", dist.Stable(1.5, 0.1, b, 0.0), obs=a) g = pyro.sample("g", dist.Stable(1.5, zero, b, a).to_event(1)) h = pyro.sample("h", dist.Stable(1.5, zero, b, 0).to_event(1), obs=a) i = pyro.sample( "i", dist.TransformedDistribution(dist.Stable(1.5, 0, b, a), dist.transforms.ExpTransform()), ) j = pyro.sample( "j", dist.TransformedDistribution(dist.Stable(1.5, 0, b, a), dist.transforms.ExpTransform()), obs=a.exp(), ) k = pyro.sample( "k", dist.TransformedDistribution(dist.Stable( 1.5, zero, b, a), dist.transforms.ExpTransform()).to_event(1), ) l = pyro.sample( "l", dist.TransformedDistribution(dist.Stable( 1.5, zero, b, a), dist.transforms.ExpTransform()).to_event(1), obs=a.exp() + zero, ) return a, b, c, d, e, f, g, h, i, j, k, l
def guide(self, bows, _): pyro.module("prop_inference_net", self.prop_inference_net) with pyro.plate("articles", bows.shape[0]): prop_mu, prop_sigma = self.prop_inference_net(bows) props = pyro.sample( "theta", dist.LogNormal(prop_mu, prop_sigma).to_event(1))
def model(self, bows, embeddings, article_ids): pyro.module("topic_recognition_net", self.topic_recognition_net) with pyro.plate("articles", bows.shape[0]): # instead of a Dirichlet prior, we use a log-normal distribution prop_mu = bows.new_zeros((bows.shape[0], self.nav_topics)) prop_sigma = bows.new_ones((bows.shape[0], self.nav_topics)) props = pyro.sample( "theta", dist.LogNormal(prop_mu, prop_sigma).to_event(1)) topics_mu, topics_sigma = self.topic_recognition_net(props) for batch_article_id, article_id in enumerate(article_ids): nav_embeddings = torch.tensor( embeddings[self.article_navs[article_id]], dtype=torch.float32).to(device) for article_nav_id in pyro.plate( "navs_{}".format(article_id), len(self.article_navs[article_id])): pyro.sample("nav_{}_{}".format(article_id, article_nav_id), dist.MultivariateNormal( topics_mu[batch_article_id], scale_tril=torch.diag( topics_sigma[batch_article_id])), obs=nav_embeddings[article_nav_id])
def gp_model(X, y): N = X.shape[0] # Priors for kernel parameters. alpha = pyro.sample("alpha", dist.LogNormal(0., 0.1)) rho = pyro.sample("rho", dist.LogNormal(0., 1.)) sigma = pyro.sample("sigma", dist.LogNormal(0., 1.)) # Covariance matrix. K = sq_exp_kernel_matrix(X, alpha, rho) + torch.eye(N) * sigma * sigma L = torch.cholesky(K) # Marginal likelihood. pyro.sample('obs', dist.MultivariateNormal(torch.zeros(N), scale_tril=L), obs=y)
def gpc(X, y): N = y.shape[0] # Priors. alpha = pyro.sample('alpha', dist.LogNormal(0, 1)) rho = pyro.sample('rho', dist.LogNormal(0, 1)) beta = pyro.sample('beta', dist.Normal(0, 1)) with pyro.plate('latent_response', N): eta = pyro.sample('eta', dist.Normal(0, 1)) # Latent function. f = compute_f(alpha, rho, beta, eta, X) with pyro.plate('response', N): pyro.sample('obs', dist.Bernoulli(logits=f), obs=y)
def model(self, X): _id = self._id N, D = X.shape global_shrinkage_prior_scale_init = self.param_init[f'global_shrinkage_prior_scale_init_{_id}'] cov_diag_prior_loc_init = self.param_init[f'cov_diag_prior_loc_init_{_id}'] cov_diag_prior_scale_init = self.param_init[f'cov_diag_prior_scale_init_{_id}'] global_shrinkage_prior_scale = pyro.param(f'global_shrinkage_scale_prior_{_id}', global_shrinkage_prior_scale_init, constraint=constraints.positive) tau = pyro.sample(f'global_shrinkage_{_id}', dist.HalfNormal(global_shrinkage_prior_scale)) b = pyro.sample('b', dist.InverseGamma(0.5,1./torch.ones(D)**2).to_event(1)) lambdasquared = pyro.sample(f'local_shrinkage_{_id}', dist.InverseGamma(0.5,1./b).to_event(1)) cov_diag_loc = pyro.param(f'cov_diag_prior_loc_{_id}', cov_diag_prior_loc_init) cov_diag_scale = pyro.param(f'cov_diag_prior_scale_{_id}', cov_diag_prior_scale_init, constraint=constraints.positive) cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale).to_event(1)) #cov_diag = cov_diag*torch.ones(D) cov_diag = cov_diag + jitter lambdasquared = lambdasquared.squeeze() if lambdasquared.dim() == 1: # outer product cov_factor_scale = torch.ger(torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,))) else: # batch outer product cov_factor_scale = torch.einsum('bp, br->bpr', torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,))) cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(0., cov_factor_scale).to_event(2)) cov_factor = cov_factor.transpose(-2,-1) with pyro.plate(f'N_{_id}', size=N, subsample_size=self.batch_size, dim=-1) as ind: X = pyro.sample('obs', dist.LowRankMultivariateNormal(torch.zeros(D), cov_factor=cov_factor, cov_diag=cov_diag), obs=X.index_select(0, ind)) return X
def factorGuide(X, batch_size, variational_parameter_initialization): N, D = X.shape K, locloc, locscale, scaleloc, scalescale, cov_factor_loc_init, cov_factor_scale_init = variational_parameter_initialization[1] with pyro.plate('D', D, dim=-1): cov_diag_loc = pyro.param('scale_loc', scaleloc) cov_diag_scale = pyro.param('scale_scale', scalescale, constraint=constraints.positive) cov_diag = pyro.sample('scale', dist.LogNormal(cov_diag_loc, cov_diag_scale)) cov_diag = cov_diag*torch.ones(D) loc_loc = pyro.param('loc_loc', locloc) loc_scale = pyro.param('loc_scale', locscale, constraint=constraints.positive) # sample variables loc = pyro.sample('loc', dist.Normal(loc_loc,loc_scale)) cov_factor = None if K > 1: with pyro.plate('K', K-1, dim=-2): cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init[:K-1,:]) cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init[:K-1,:], constraint=constraints.positive) cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc, cov_factor_scale)) cov_factor_new_loc = pyro.param('cov_factor_new_loc_{}'.format(K), cov_factor_loc_init[-1,:]) cov_factor_new_scale = pyro.param('cov_factor_new_scale_{}'.format(K), cov_factor_scale_init[-1,:], constraint=constraints.positive) cov_factor_new = pyro.sample('cov_factor_new', dist.Normal(cov_factor_new_loc,cov_factor_new_scale)) cov_factor = torch.cat([cov_factor, torch.unsqueeze(cov_factor_new, dim=0)]) else: with pyro.plate('K', K): cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init) cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive) cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale)) cov_factor = cov_factor.transpose(0,1) return loc, cov_factor, cov_diag
def model(self, zero_data, covariates): with pyro.plate("batch", len(zero_data), dim=-2): zero_data = pyro.subsample(zero_data, event_dim=1) covariates = pyro.subsample(covariates, event_dim=1) loc = zero_data[..., :1, :] scale = pyro.sample("scale", dist.LogNormal(loc, 1).to_event(1)) with self.time_plate: jumps = pyro.sample("jumps", dist.Normal(0, scale).to_event(1)) prediction = jumps.cumsum(-2) duration, obs_dim = zero_data.shape[-2:] noise_dist = dist.LinearHMM( dist.Stable(1.9, 0).expand([obs_dim]).to_event(1), torch.eye(obs_dim), dist.Stable(1.9, 0).expand([obs_dim]).to_event(1), torch.eye(obs_dim), dist.Stable(1.9, 0).expand([obs_dim]).to_event(1), duration=duration, ) rep = StableReparam() with poutine.reparam( config={"residual": LinearHMMReparam(rep, rep, rep)}): self.predict(noise_dist, prediction)
def projectedMixtureGuide(X, batch_size, variational_parameter_initialization): """ Covariances of all clusters are locked, we're just learning mixture weights and means """ N, D = X.shape loc_loc, loc_scale, scale_loc, scale_scale, component_logits_concentration, cov_factor_loc_init,cov_factor_scale_init = variational_parameter_initialization[1] C = loc_loc.shape[0] K = cov_factor_loc_init.shape[0] component_logits_concentration = pyro.param('component_logits_concentration', component_logits_concentration, constraint=constraints.positive) component_logits = pyro.sample('component_logits', dist.Dirichlet(component_logits_concentration)) with pyro.plate('D', D): cov_diag_loc = pyro.param('scale_loc', scale_loc) cov_diag_scale = pyro.param('scale_scale', scale_scale, constraint=constraints.positive) cov_diag = pyro.sample('scale', dist.LogNormal(scale_loc, scale_scale)) cov_diag = cov_diag*torch.ones(D) with pyro.plate('K', K): cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init) cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive) cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale)) cov_factor = cov_factor.transpose(0,1) with pyro.plate('C', C): loc_loc = pyro.param('loc_loc', loc_loc) loc_scale = pyro.param('loc_scale', loc_scale, constraint=constraints.positive) locs = pyro.sample('locs', dist.Normal(loc_loc,loc_scale)) return component_logits, cov_diag, cov_factor, locs
def zeroMeanFactorGuide(X, batch_size, variational_parameter_initialization): N, D = X.shape K, scaleloc, scalescale, cov_factor_loc_init, cov_factor_scale_init = variational_parameter_initialization[1] with pyro.plate('D', D, dim=-1): cov_diag_loc = pyro.param('scale_loc', scaleloc, constraint=constraints.positive) cov_diag_scale = pyro.param('scale_scale', scalescale, constraint=constraints.positive) cov_diag = pyro.sample('scale', dist.LogNormal(cov_diag_loc, cov_diag_scale)) cov_diag = cov_diag*torch.ones(D) # sample variables cov_factor = None if K > 1: with pyro.plate('K', K-1, dim=-2): cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init[:K-1,:]) cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init[:K-1,:], constraint=constraints.positive) cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc, cov_factor_scale)) cov_factor_new_loc = pyro.param('cov_factor_new_loc_{}'.format(K), cov_factor_loc_init[-1,:]) cov_factor_new_scale = pyro.param('cov_factor_new_scale_{}'.format(K), cov_factor_scale_init[-1,:], constraint=constraints.positive) cov_factor_new = pyro.sample('cov_factor_new', dist.Normal(cov_factor_new_loc,cov_factor_new_scale)) # when using pyro.infer.Predictive, cov_factor_new is somehow sampled as 2-d tensors instead of 1-d if cov_factor_new.dim() == cov_factor.dim(): cov_factor = torch.cat([cov_factor, cov_factor_new], dim=-2) else: cov_factor = torch.cat([cov_factor, torch.unsqueeze(cov_factor_new, dim=-2)], dim=-2) else: with pyro.plate('K', K): cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init) cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive) cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale)) cov_factor = cov_factor.transpose(-2,-1) return cov_factor, cov_diag
def sphericalMixture(X, batch_size, prior_parameters): """ Covariances of all clusters are diagonal, only params are mixture weights, means and shared variable variances ...so it's actually an ellipsoidal mixture, if you wanna be anal about it. C**t. """ N, D = X.shape locloc, locscale, scaleloc, scalescale, component_logits_concentration = prior_parameters[0] # get number of clusters from first dimension of means C = locloc.shape[0] component_logits = pyro.sample('component_logits', dist.Dirichlet(component_logits_concentration)) with pyro.plate('D', D): cov_diag = pyro.sample('scale', dist.LogNormal(scaleloc, scalescale)) with pyro.plate('C', C): locs = pyro.sample('locs', dist.Normal(locloc,locscale)) with pyro.plate('N', size=N, subsample_size=batch_size) as ind: assignment = pyro.sample('assignment', dist.Categorical(component_logits), infer={"enumerate": "parallel"}) #X = pyro.sample('obs', dist.MultivariateNormal(locs.index_select(-2, assignment), torch.diag(cov_diag)), obs=X.index_select(0, ind)) # use index_select instead of locs[assignment] so batching works correctly # have to select the right index both in a parallel enumeration context and a batch context # leading to this ugly hack indexed_locs = locs.index_select(-2, assignment.squeeze()).view(*locs.shape[:-2],*assignment.shape,locs.shape[-1]) #print(locs.shape) #print(assignment.shape) #print(indexed_locs.shape) X = pyro.sample('obs', dist.MultivariateNormal(indexed_locs, torch.diag_embed(cov_diag)), obs=X.index_select(0, ind)) return X
def test_transformed_hmm_shape(batch_shape, duration, hidden_dim, obs_dim): init_dist = random_mvn(batch_shape, hidden_dim) trans_mat = torch.randn(batch_shape + (duration, hidden_dim, hidden_dim)) trans_dist = random_mvn(batch_shape + (duration, ), hidden_dim) obs_mat = torch.randn(batch_shape + (duration, hidden_dim, obs_dim)) obs_dist = dist.LogNormal( torch.randn(batch_shape + (duration, obs_dim)), torch.rand(batch_shape + (duration, obs_dim)).exp()).to_event(1) hmm = dist.LinearHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=duration) def model(data=None): with pyro.plate_stack("plates", batch_shape): return pyro.sample("x", hmm, obs=data) data = model() with poutine.trace() as tr: with poutine.reparam(config={"x": LinearHMMReparam()}): model(data) fn = tr.trace.nodes["x"]["fn"] assert isinstance(fn, dist.TransformedDistribution) assert isinstance(fn.base_dist, dist.GaussianHMM) tr.trace.compute_log_prob() # smoke test only
def model(self, x, y=None): # Register various nn.Modules with Pyro pyro.module("scanvi", self) # This gene-level parameter modulates the variance of the observation distribution theta = pyro.param("inverse_dispersion", 10.0 * x.new_ones(self.num_genes), constraint=constraints.positive) # We scale all sample statements by scale_factor so that the ELBO is normalized # wrt the number of datapoints and genes with pyro.plate("batch", len(x)), poutine.scale(scale=self.scale_factor): z1 = pyro.sample("z1", dist.Normal(0, x.new_ones(self.latent_dim)).to_event(1)) # Note that if y is None (i.e. y is unobserved) then y will be sampled; # otherwise y will be treated as observed. y = pyro.sample("y", dist.OneHotCategorical(logits=x.new_zeros(self.num_labels)), obs=y) z2_loc, z2_scale = self.z2_decoder(z1, y) z2 = pyro.sample("z2", dist.Normal(z2_loc, z2_scale).to_event(1)) l_scale = self.l_scale * x.new_ones(1) l = pyro.sample("l", dist.LogNormal(self.l_loc, l_scale).to_event(1)) # Note that by construction mu is normalized (i.e. mu.sum(-1) == 1) and the # total scale of counts for each cell is determined by `l` gate_logits, mu = self.x_decoder(z2) # TODO revisit this parameterization if torch.distributions.NegativeBinomial changes # from failure to success parametrization; # see https://github.com/pytorch/pytorch/issues/42449 nb_logits = (l * mu + self.epsilon).log() - (theta + self.epsilon).log() x_dist = dist.ZeroInflatedNegativeBinomial(gate_logits=gate_logits, total_count=theta, logits=nb_logits) # Observe the datapoint x using the observation distribution x_dist pyro.sample("x", x_dist.to_event(1), obs=x)
def setUp(self): self.mu = Variable(torch.Tensor([1.4])) self.sigma = Variable(torch.Tensor([0.4])) self.test_data = Variable(torch.Tensor([5.5])) self.batch_mu = Variable(torch.Tensor([[1.4], [2.6]])) self.batch_sigma = Variable(torch.Tensor([[0.4], [0.5]])) self.batch_test_data = Variable(torch.Tensor([[5.5], [6.4]])) self.dist = dist.LogNormal(self.mu, self.sigma) self.batch_dist = dist.LogNormal(self.batch_mu, self.batch_sigma) self.analytic_mean = torch.exp( self.mu + 0.5 * torch.pow(self.sigma, 2.0)).data.cpu().numpy()[0] var_factor = torch.exp(torch.pow(self.sigma, 2.0)) - Variable( torch.ones(1)) self.analytic_var = var_factor.data.cpu().numpy()[0] * np.square( self.analytic_mean) self.n_samples = 70000