def run(
task: Task,
num_samples: int,
num_simulations: int,
num_observation: Optional[int] = None,
observation: Optional[torch.Tensor] = None,
num_rounds: int = 10,
neural_net: str = "resnet",
hidden_features: int = 50,
simulation_batch_size: int = 1000,
training_batch_size: int = 10000,
num_atoms: int = 10,
automatic_transforms_enabled: bool = True,
mcmc_method: str = "slice_np_vectorized",
mcmc_parameters: Dict[str, Any] = {
"num_chains": 100,
"thin": 10,
"warmup_steps": 100,
"init_strategy": "sir",
"sir_batch_size": 1000,
"sir_num_batches": 100,
},
z_score_x: bool = True,
z_score_theta: bool = True,
variant: str = "B",
) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
"""Runs (S)NRE from `sbi`
Args:
task: Task instance
num_samples: Number of samples to generate from posterior
num_observation: Observation number to load, alternative to `observation`
observation: Observation, alternative to `num_observation`
num_simulations: Simulation budget
num_rounds: Number of rounds
neural_net: Neural network to use, one of linear / mlp / resnet
hidden_features: Number of hidden features in network
simulation_batch_size: Batch size for simulator
training_batch_size: Batch size for training network
num_atoms: Number of atoms, -1 means same as `training_batch_size`
automatic_transforms_enabled: Whether to enable automatic transforms
mcmc_method: MCMC method
mcmc_parameters: MCMC parameters
z_score_x: Whether to z-score x
z_score_theta: Whether to z-score theta
variant: Can be used to switch between SNRE-A (AALR) and -B (SRE)
Returns:
Samples from posterior, number of simulator calls, log probability of true params if computable
"""
assert not (num_observation is None and observation is None)
assert not (num_observation is not None and observation is not None)
log = logging.getLogger(__name__)
if num_rounds == 1:
log.info(f"Running NRE")
num_simulations_per_round = num_simulations
else:
log.info(f"Running SNRE")
num_simulations_per_round = math.floor(num_simulations / num_rounds)
if simulation_batch_size > num_simulations_per_round:
simulation_batch_size = num_simulations_per_round
log.warn("Reduced simulation_batch_size to num_simulation_per_round")
if training_batch_size > num_simulations_per_round:
training_batch_size = num_simulations_per_round
log.warn("Reduced training_batch_size to num_simulation_per_round")
prior = task.get_prior_dist()
if observation is None:
observation = task.get_observation(num_observation)
simulator = task.get_simulator(max_calls=num_simulations)
transforms = task._get_transforms(
automatic_transforms_enabled)["parameters"]
if automatic_transforms_enabled:
prior = wrap_prior_dist(prior, transforms)
simulator = wrap_simulator_fn(simulator, transforms)
classifier = classifier_nn(
model=neural_net.lower(),
hidden_features=hidden_features,
z_score_x=z_score_x,
z_score_theta=z_score_theta,
)
if variant == "A":
inference_class = inference.SNRE_A
inference_method_kwargs = {}
elif variant == "B":
inference_class = inference.SNRE_B
inference_method_kwargs = {"num_atoms": num_atoms}
else:
raise NotImplementedError
inference_method = inference_class(classifier=classifier, prior=prior)
posteriors = []
proposal = prior
mcmc_parameters["warmup_steps"] = 25
for r in range(num_rounds):
theta, x = inference.simulate_for_sbi(
simulator,
proposal,
num_simulations=num_simulations_per_round,
simulation_batch_size=simulation_batch_size,
)
density_estimator = inference_method.append_simulations(
theta, x, from_round=r).train(
training_batch_size=training_batch_size,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
show_train_summary=True,
**inference_method_kwargs,
)
if r > 1:
mcmc_parameters["init_strategy"] = "latest_sample"
posterior = inference_method.build_posterior(
density_estimator,
mcmc_method=mcmc_method,
mcmc_parameters=mcmc_parameters)
# Copy hyperparameters, e.g., mcmc_init_samples for "latest_sample" strategy.
if r > 0:
posterior.copy_hyperparameters_from(posteriors[-1])
proposal = posterior.set_default_x(observation)
posteriors.append(posterior)
posterior = wrap_posterior(posteriors[-1], transforms)
assert simulator.num_simulations == num_simulations
samples = posterior.sample((num_samples, )).detach()
return samples, simulator.num_simulations, None
def run(
task: Task,
num_samples: int,
num_simulations: int,
num_observation: Optional[int] = None,
observation: Optional[torch.Tensor] = None,
num_rounds: int = 10,
neural_net: str = "nsf",
hidden_features: int = 50,
simulation_batch_size: int = 1000,
training_batch_size: int = 10000,
num_atoms: int = 10,
automatic_transforms_enabled: bool = False,
z_score_x: bool = True,
z_score_theta: bool = True,
) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
"""Runs (S)NPE from `sbi`
Args:
task: Task instance
num_samples: Number of samples to generate from posterior
num_simulations: Simulation budget
num_observation: Observation number to load, alternative to `observation`
observation: Observation, alternative to `num_observation`
num_rounds: Number of rounds
neural_net: Neural network to use, one of maf / mdn / made / nsf
hidden_features: Number of hidden features in network
simulation_batch_size: Batch size for simulator
training_batch_size: Batch size for training network
num_atoms: Number of atoms, -1 means same as `training_batch_size`
automatic_transforms_enabled: Whether to enable automatic transforms
z_score_x: Whether to z-score x
z_score_theta: Whether to z-score theta
Returns:
Samples from posterior, number of simulator calls, log probability of true params if computable
"""
assert not (num_observation is None and observation is None)
assert not (num_observation is not None and observation is not None)
log = logging.getLogger(__name__)
if num_rounds == 1:
log.info(f"Running NPE")
num_simulations_per_round = num_simulations
else:
log.info(f"Running SNPE")
num_simulations_per_round = math.floor(num_simulations / num_rounds)
if simulation_batch_size > num_simulations_per_round:
simulation_batch_size = num_simulations_per_round
log.warn("Reduced simulation_batch_size to num_simulation_per_round")
if training_batch_size > num_simulations_per_round:
training_batch_size = num_simulations_per_round
log.warn("Reduced training_batch_size to num_simulation_per_round")
prior = task.get_prior_dist()
if observation is None:
observation = task.get_observation(num_observation)
simulator = task.get_simulator(max_calls=num_simulations)
transforms = task._get_transforms(
automatic_transforms_enabled)["parameters"]
if automatic_transforms_enabled:
prior = wrap_prior_dist(prior, transforms)
simulator = wrap_simulator_fn(simulator, transforms)
density_estimator_fun = posterior_nn(
model=neural_net.lower(),
hidden_features=hidden_features,
z_score_x=z_score_x,
z_score_theta=z_score_theta,
)
inference_method = inference.SNPE_C(
prior, density_estimator=density_estimator_fun)
posteriors = []
proposal = prior
for _ in range(num_rounds):
theta, x = inference.simulate_for_sbi(
simulator,
proposal,
num_simulations=num_simulations_per_round,
simulation_batch_size=simulation_batch_size,
)
density_estimator = inference_method.append_simulations(
theta, x, proposal=proposal).train(
num_atoms=num_atoms,
training_batch_size=training_batch_size,
retrain_from_scratch_each_round=False,
discard_prior_samples=False,
use_combined_loss=False,
show_train_summary=True,
)
posterior = inference_method.build_posterior(density_estimator,
sample_with_mcmc=False)
proposal = posterior.set_default_x(observation)
posteriors.append(posterior)
posterior = wrap_posterior(posteriors[-1], transforms)
assert simulator.num_simulations == num_simulations
samples = posterior.sample((num_samples, )).detach()
if num_observation is not None:
true_parameters = task.get_true_parameters(
num_observation=num_observation)
log_prob_true_parameters = posterior.log_prob(true_parameters)
return samples, simulator.num_simulations, log_prob_true_parameters
else:
return samples, simulator.num_simulations, None
def run(
task: Task,
num_samples: int,
num_simulations: int,
num_simulations_per_step: int = 100,
num_observation: Optional[int] = None,
observation: Optional[torch.Tensor] = None,
automatic_transforms_enabled: bool = False,
mcmc_method: str = "slice_np",
mcmc_parameters: Dict[str, Any] = {},
diag_eps: float = 0.0,
) -> (torch.Tensor, int, Optional[torch.Tensor]):
"""Runs (S)NLE from `sbi`
Args:
task: Task instance
num_observation: Observation number to load, alternative to `observation`
observation: Observation, alternative to `num_observation`
num_samples: Number of samples to generate from posterior
num_simulations: Simulation budget
num_simulations_per_step: Number of simulations per MCMC step
automatic_transforms_enabled: Whether to enable automatic transforms
mcmc_method: MCMC method
mcmc_parameters: MCMC parameters
diag_eps: Epsilon applied to diagonal
Returns:
Samples from posterior, number of simulator calls, log probability of true params if computable
"""
assert not (num_observation is None and observation is None)
assert not (num_observation is not None and observation is not None)
log = logging.getLogger(__name__)
log.info(f"Running SL")
prior = task.get_prior_dist()
if observation is None:
observation = task.get_observation(num_observation)
simulator = task.get_simulator()
transforms = task._get_transforms(automatic_transforms_enabled)["parameters"]
prior = wrap_prior_dist(prior, transforms)
simulator = wrap_simulator_fn(simulator, transforms)
likelihood_estimator = SynthLikNet(
simulator=simulator,
num_simulations_per_step=num_simulations_per_step,
diag_eps=diag_eps,
)
posterior = LikelihoodBasedPosterior(
method_family="snle",
neural_net=likelihood_estimator,
prior=prior,
x_shape=observation.shape,
mcmc_parameters=mcmc_parameters,
)
posterior.set_default_x(observation)
posterior = wrap_posterior(posterior, transforms)
# assert simulator.num_simulations == num_simulations
samples = posterior.sample((num_samples,)).detach()
return samples, simulator.num_simulations, None