def test_c2st_and_map_snl_on_linearGaussian_different(num_dim: int, prior_str: str): """Test SNL on linear Gaussian, comparing to ground truth posterior via c2st. Args: num_dim: parameter dimension of the gaussian model prior_str: one of "gaussian" or "uniform" """ num_samples = 500 num_simulations = 3000 trials_to_test = [1] # likelihood_mean will be likelihood_shift+theta likelihood_shift = -1.0 * ones(num_dim) # Use increased cov to avoid too small posterior cov for many trials. likelihood_cov = 0.8 * eye(num_dim) if prior_str == "gaussian": prior_mean = zeros(num_dim) prior_cov = eye(num_dim) prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov) else: prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim)) simulator, prior = prepare_for_sbi( lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior) density_estimator = likelihood_nn("maf", num_transforms=3) inference = SNLE(density_estimator=density_estimator, show_progress_bars=False) theta, x = simulate_for_sbi(simulator, prior, num_simulations, simulation_batch_size=10000) likelihood_estimator = inference.append_simulations(theta, x).train() # Test inference amortized over trials. for num_trials in trials_to_test: x_o = zeros((num_trials, num_dim)) if prior_str == "gaussian": gt_posterior = true_posterior_linear_gaussian_mvn_prior( x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov) target_samples = gt_posterior.sample((num_samples, )) elif prior_str == "uniform": target_samples = samples_true_posterior_linear_gaussian_uniform_prior( x_o, likelihood_shift, likelihood_cov, prior=prior, num_samples=num_samples, ) else: raise ValueError(f"Wrong prior_str: '{prior_str}'.") potential_fn, theta_transform = likelihood_estimator_based_potential( prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o) posterior = MCMCPosterior( proposal=prior, potential_fn=potential_fn, theta_transform=theta_transform, method="slice_np_vectorized", thin=5, num_chains=5, ) samples = posterior.sample(sample_shape=(num_samples, )) # Check performance based on c2st accuracy. check_c2st(samples, target_samples, alg=f"snle_a-{prior_str}-prior-{num_trials}-trials") map_ = posterior.map(num_init_samples=1_000, init_method="proposal", show_progress_bars=False) # TODO: we do not have a test for SNL log_prob(). This is because the output # TODO: density is not normalized, so KLd does not make sense. if prior_str == "uniform": # Check whether the returned probability outside of the support is zero. posterior_prob = get_prob_outside_uniform_prior( posterior, prior, num_dim) assert ( posterior_prob == 0.0 ), "The posterior probability outside of the prior support is not zero" assert ((map_ - ones(num_dim))**2).sum() < 0.5 else: assert ((map_ - gt_posterior.mean)**2).sum() < 0.5
def test_c2st_sre_variants_on_linearGaussian( num_dim: int, num_trials: int, prior_str: str, method_str: str, ): """Test c2st accuracy of inference with SRE on linear Gaussian model. Args: num_dim: parameter dimension of the gaussian model prior_str: one of "gaussian" or "uniform" """ x_o = zeros(num_trials, num_dim) num_samples = 500 num_simulations = 2600 if num_trials == 1 else 40500 # `likelihood_mean` will be `likelihood_shift + theta`. likelihood_shift = -1.0 * ones(num_dim) likelihood_cov = 0.8 * eye(num_dim) if prior_str == "gaussian": prior_mean = zeros(num_dim) prior_cov = eye(num_dim) prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov) else: prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim)) def simulator(theta): return linear_gaussian(theta, likelihood_shift, likelihood_cov) simulator, prior = prepare_for_sbi(simulator, prior) kwargs = dict( classifier="resnet", show_progress_bars=False, ) inference = SNRE_B(**kwargs) if method_str == "sre" else AALR(**kwargs) # Should use default `num_atoms=10` for SRE; `num_atoms=2` for AALR theta, x = simulate_for_sbi( simulator, prior, num_simulations, simulation_batch_size=50 ) ratio_estimator = inference.append_simulations(theta, x).train() potential_fn, theta_transform = ratio_estimator_based_potential( ratio_estimator=ratio_estimator, prior=prior, x_o=x_o ) posterior = MCMCPosterior( potential_fn=potential_fn, theta_transform=theta_transform, proposal=prior, method="slice_np_vectorized", thin=5, num_chains=5, ) samples = posterior.sample(sample_shape=(num_samples,)) # Get posterior samples. if prior_str == "gaussian": gt_posterior = true_posterior_linear_gaussian_mvn_prior( x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov ) target_samples = gt_posterior.sample((num_samples,)) else: target_samples = samples_true_posterior_linear_gaussian_uniform_prior( x_o, likelihood_shift, likelihood_cov, prior=prior, num_samples=num_samples ) # Check performance based on c2st accuracy. check_c2st( samples, target_samples, alg=f"sre-{prior_str}-{method_str}-{num_trials}trials" ) map_ = posterior.map(num_init_samples=1_000, init_method="proposal") # Checks for log_prob() if prior_str == "gaussian" and method_str == "aalr": # For the Gaussian prior, we compute the KLd between ground truth and # posterior. We can do this only if the classifier_loss was as described in # Hermans et al. 2020 ('aalr') since Durkan et al. 2020 version only allows # evaluation up to a constant. # For the Gaussian prior, we compute the KLd between ground truth and posterior dkl = get_dkl_gaussian_prior( posterior, x_o, likelihood_shift, likelihood_cov, prior_mean, prior_cov ) max_dkl = 0.15 assert ( dkl < max_dkl ), f"KLd={dkl} is more than 2 stds above the average performance." assert ((map_ - gt_posterior.mean) ** 2).sum() < 0.5 if prior_str == "uniform": # Check whether the returned probability outside of the support is zero. posterior_prob = get_prob_outside_uniform_prior(posterior, prior, num_dim) assert ( posterior_prob == 0.0 ), "The posterior probability outside of the prior support is not zero" assert ((map_ - ones(num_dim)) ** 2).sum() < 0.5