def sample_nuts(self, num_results=int(10e3), n_burnin=int(5e3), max_tree_depth=10, step_size=0.01, num_adapt_steps=4000): """ NUTS sampling - WIP, DO NOT USE!!! Parameters ---------- num_results -- int MCMC chain length (default 20000) n_burnin -- int Number of burnin iterations (default 5000) max_tre_depth -- int Maximum tree depth (default 10) step_size -- float Initial step size (default 0.01) Returns ------- error NotImplementedError """ constraining_bijectors = [ tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), ] # NUTS transition kernel nuts_kernel = tfp.mcmc.NoUTurnSampler( target_log_prob_fn=self.target_log_prob_fn, step_size=step_size, max_tree_depth=max_tree_depth) nuts_kernel = tfp.mcmc.TransformedTransitionKernel( inner_kernel=nuts_kernel, bijector=constraining_bijectors) nuts_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation( inner_kernel=nuts_kernel, num_adaptation_steps=num_adapt_steps, target_accept_prob=0.8, step_size_setter_fn=lambda pkr, new_step_size: pkr._replace( inner_results=pkr.inner_results._replace(step_size= new_step_size)), step_size_getter_fn=lambda pkr: pkr.inner_results.step_size, log_accept_prob_getter_fn=lambda pkr: pkr.inner_results. log_accept_ratio, ) # HMC sampling states, kernel_results, duration = self.sampling(num_results, n_burnin, nuts_kernel, self.params, is_nuts=True) states_burnin, acc_rate = self.get_chains_after_burnin(states, kernel_results, n_burnin, is_nuts=True) y_hat = self.get_y_hat(states_burnin, num_results, n_burnin) params = dict(zip(self.param_names, states_burnin)) # Result object generation setup if self.baseline_index is not None: cell_types_nb = self.cell_types[:self. baseline_index] + self.cell_types[ self.baseline_index + 1:] else: cell_types_nb = self.cell_types posterior = { var_name: [var] for var_name, var in params.items() if "prediction" not in var_name } posterior_predictive = {"prediction": [params["prediction"]]} observed_data = {"y": self.y} dims = { "alpha": ["cell_type"], "mu_b": ["1"], "sigma_b": ["1"], "b_offset": ["covariate", "cell_type_nb"], "ind_raw": ["covariate", "cell_type_nb"], "ind": ["covariate", "cell_type_nb"], "b_raw": ["covariate", "cell_type_nb"], "beta": ["covariate", "cell_type"], "concentration": ["sample", "cell_type"], "prediction": ["sample", "cell_type"] } coords = { "cell_type": self.cell_types, "cell_type_nb": cell_types_nb, "covariate": self.covariate_names, "sample": range(self.y.shape[0]) } sampling_stats = { "chain_length": num_results, "n_burnin": n_burnin, "acc_rate": acc_rate, "duration": duration, "y_hat": y_hat } model_specs = { "baseline": self.baseline_index, "formula": self.formula } return res.CAResultConverter(posterior=posterior, posterior_predictive=posterior_predictive, observed_data=observed_data, dims=dims, coords=coords).to_result_data( sampling_stats=sampling_stats, model_specs=model_specs)
def sample_hmc(self, num_results=int(20e3), n_burnin=int(5e3), num_leapfrog_steps=10, step_size=0.01): """ HMC sampling Parameters ---------- num_results -- int MCMC chain length (default 20000) n_burnin -- int Number of burnin iterations (default 5000) num_leapfrog_steps -- int HMC leapfrog steps (default 10) step_size -- float Initial step size (default 0.01) Returns ------- result scdcdm.util.result_data object """ # (not in use atm) constraining_bijectors = [ tfb.Identity(), tfb.Identity(), ] # HMC transition kernel hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo( target_log_prob_fn=self.target_log_prob_fn, step_size=step_size, num_leapfrog_steps=num_leapfrog_steps) hmc_kernel = tfp.mcmc.TransformedTransitionKernel( inner_kernel=hmc_kernel, bijector=constraining_bijectors) hmc_kernel = tfp.mcmc.SimpleStepSizeAdaptation( inner_kernel=hmc_kernel, num_adaptation_steps=int(4000), target_accept_prob=0.8) # HMC sampling states, kernel_results = self.sampling(num_results, n_burnin, hmc_kernel, self.params) states_burnin = self.get_chains_after_burnin(states, kernel_results, n_burnin) y_hat = self.get_y_hat(states_burnin, num_results, n_burnin) params = dict(zip(self.param_names, states_burnin)) posterior = { var_name: [var] for var_name, var in params.items() if "prediction" not in var_name } posterior_predictive = {"prediction": [params["prediction"]]} observed_data = {"y": self.y} dims = { "alpha": ["cell_type"], "beta": ["covariate", "cell_type"], "concentration": ["sample", "cell_type"], "prediction": ["sample", "cell_type"] } coords = { "cell_type": self.cell_types, "covariate": self.covariate_names, "sample": range(self.y.shape[0]) } return res.CAResultConverter(posterior=posterior, posterior_predictive=posterior_predictive, observed_data=observed_data, dims=dims, coords=coords).to_result_data( y_hat, baseline=False)
def sample_hmc_da(self, num_results=int(20e3), n_burnin=int(5e3), num_leapfrog_steps=10, step_size=0.01, num_adapt_steps=4000): """ HMC sampling Parameters ---------- num_results -- int MCMC chain length (default 20000) n_burnin -- int Number of burnin iterations (default 5000) num_leapfrog_steps -- int HMC leapfrog steps (default 10) step_size -- float Initial step size (default 0.01) Returns ------- result -- scdcdm.util.result_data.CAResult object Compositional analysis result """ # (not in use atm) constraining_bijectors = [ tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), ] # HMC transition kernel hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo( target_log_prob_fn=self.target_log_prob_fn, step_size=step_size, num_leapfrog_steps=num_leapfrog_steps) hmc_kernel = tfp.mcmc.TransformedTransitionKernel( inner_kernel=hmc_kernel, bijector=constraining_bijectors) # if n_burnin == 0: # adapt_steps = int(0.2*num_results) # else: # adapt_steps = int(0.8*n_burnin) hmc_kernel = tfp.mcmc.DualAveragingStepSizeAdaptation( inner_kernel=hmc_kernel, num_adaptation_steps=num_adapt_steps, target_accept_prob=0.8) # HMC sampling states, kernel_results, duration = self.sampling( num_results, n_burnin, hmc_kernel, self.params) states_burnin, acc_rate = self.get_chains_after_burnin( states, kernel_results, n_burnin) y_hat = self.get_y_hat(states_burnin, num_results, n_burnin) params = dict(zip(self.param_names, states_burnin)) # Result object generation setup if self.baseline_index is not None: cell_types_nb = self.cell_types[:self. baseline_index] + self.cell_types[ self.baseline_index + 1:] else: cell_types_nb = self.cell_types posterior = { var_name: [var] for var_name, var in params.items() if "prediction" not in var_name } posterior_predictive = {"prediction": [params["prediction"]]} observed_data = {"y": self.y} dims = { "alpha": ["cell_type"], "mu_b": ["1"], "sigma_b": ["1"], "b_offset": ["covariate", "cell_type_nb"], "ind_raw": ["covariate", "cell_type_nb"], "ind": ["covariate", "cell_type_nb"], "b_raw": ["covariate", "cell_type_nb"], "beta": ["covariate", "cell_type"], "concentration": ["sample", "cell_type"], "prediction": ["sample", "cell_type"] } coords = { "cell_type": self.cell_types, "cell_type_nb": cell_types_nb, "covariate": self.covariate_names, "sample": range(self.y.shape[0]) } sampling_stats = { "chain_length": num_results, "n_burnin": n_burnin, "acc_rate": acc_rate, "duration": duration, "y_hat": y_hat } model_specs = { "baseline": self.baseline_index, "formula": self.formula } return res.CAResultConverter(posterior=posterior, posterior_predictive=posterior_predictive, observed_data=observed_data, dims=dims, coords=coords).to_result_data( sampling_stats=sampling_stats, model_specs=model_specs)
def sample_nuts(self, num_results=int(10e3), n_burnin=int(5e3), max_tree_depth=10, step_size=0.01): """ NUTS sampling - WIP, DO NOT USE!!! Parameters ---------- num_results -- int MCMC chain length (default 20000) n_burnin -- int Number of burnin iterations (default 5000) max_tre_depth -- int Maximum tree depth (default 10) step_size -- float Initial step size (default 0.01) Returns ------- error NotImplementedError """ raise NotImplementedError # (not in use atm) constraining_bijectors = [ tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), tfb.Identity(), ] # NUTS transition kernel nuts_kernel = tfp.mcmc.NoUTurnSampler( target_log_prob_fn=self.target_log_prob_fn, step_size=step_size, max_tree_depth=max_tree_depth) nuts_kernel = tfp.mcmc.TransformedTransitionKernel( inner_kernel=nuts_kernel, bijector=constraining_bijectors) nuts_kernel = tfp.mcmc.SimpleStepSizeAdaptation( inner_kernel=nuts_kernel, num_adaptation_steps=int(4000), target_accept_prob=0.9, step_size_setter_fn=lambda pkr, new_step_size: pkr._replace( inner_results=pkr.inner_results._replace(step_size= new_step_size)), step_size_getter_fn=lambda pkr: pkr.inner_results.step_size, log_accept_prob_getter_fn=lambda pkr: pkr.inner_results. log_accept_ratio, ) states, kernel_results = self.sampling(num_results, n_burnin, nuts_kernel, self.params) states_burnin = self.get_chains_after_burnin(states, kernel_results, n_burnin) y_hat = self.get_y_hat(states_burnin, num_results, n_burnin) params = dict(zip(self.param_names, states_burnin)) # Arviz setup posterior = { var_name: [var] for var_name, var in params.items() if "prediction" not in var_name } posterior_predictive = {"prediction": [params["prediction"]]} observed_data = {"y": self.y} dims = { "alpha": ["cell_type"], "mu_b": ["1"], "sigma_b": ["1"], "b_offset": ["covariate", "cell_type"], "ind_raw": ["covariate", "cell_type"], "ind": ["covariate", "cell_type"], "b_raw": ["covariate", "cell_type"], "beta": ["cov", "cell_type"], "concentration": ["sample", "cell_type"], "prediction": ["sample", "cell_type"] } coords = { "cell_type": self.cell_types, "covariate": self.covariate_names, "sample": range(self.y.shape[0]) } return res.CAResultConverter(posterior=posterior, posterior_predictive=posterior_predictive, observed_data=observed_data, dims=dims, coords=coords).to_result_data( y_hat, baseline=False)