Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix latest sample init strategy for snle and snre. #57

Merged
merged 2 commits into from
Jun 1, 2023
Merged

Fix latest sample init strategy for snle and snre. #57

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2242860
copy init params from prev posterior, adapt tests.
janfb Jun 1, 2023
de2c6ff
adapt to new sbi z-score api kwargs.
janfb Jun 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions sbibm/algorithms/sbi/snle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,8 @@ def run(
"num_candidate_samples": 10000,
},
},
z_score_x: bool = True,
z_score_theta: bool = True,
z_score_x: str = "independent",
z_score_theta: str = "independent",
max_num_epochs: Optional[int] = 2**31 - 1,
) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
"""Runs (S)NLE from `sbi`
Expand Down Expand Up @@ -126,15 +126,19 @@ def run(
show_train_summary=True,
max_num_epochs=max_num_epochs,
)
if r > 1:
mcmc_parameters["init_strategy"] = "latest_sample"

posterior = inference_method.build_posterior(
density_estimator=density_estimator,
sample_with="mcmc",
mcmc_method=mcmc_method,
mcmc_parameters=mcmc_parameters,
)
# Change init_strategy to latest_sample after second round.
if r > 1:
posterior.init_strategy = "latest_sample"
# copy init params from round 2 posterior.
posterior._mcmc_init_params = proposal._mcmc_init_params

proposal = posterior.set_default_x(observation)
posteriors.append(posterior)

Expand Down
4 changes: 2 additions & 2 deletions sbibm/algorithms/sbi/snpe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ def run(
training_batch_size: int = 10000,
num_atoms: int = 10,
automatic_transforms_enabled: bool = False,
z_score_x: bool = True,
z_score_theta: bool = True,
z_score_x: str = "independent",
z_score_theta: str = "independent",
max_num_epochs: Optional[int] = 2**31 - 1,
) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
"""Runs (S)NPE from `sbi`
Expand Down
12 changes: 8 additions & 4 deletions sbibm/algorithms/sbi/snre.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ def run(
"num_candidate_samples": 10000,
},
},
z_score_x: bool = True,
z_score_theta: bool = True,
z_score_x: str = "independent",
z_score_theta: str = "independent",
variant: str = "B",
max_num_epochs: Optional[int] = 2**31 - 1,
) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
Expand Down Expand Up @@ -137,13 +137,17 @@ def run(
max_num_epochs=max_num_epochs,
**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,
)
# Change init_strategy to latest_sample after second round.
if r > 1:
posterior.init_strategy = "latest_sample"
# copy init params from round 2 posterior.
posterior._mcmc_init_params = proposal._mcmc_init_params
proposal = posterior.set_default_x(observation)
posteriors.append(posterior)

Expand Down
3 changes: 2 additions & 1 deletion tests/algorithms/sbi/test_sbi_run_methods.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,13 @@ def test_sbi_api(
run_method, task_name, num_observation, num_simulations=2_000, num_samples=100
):
task = sbibm.get_task(task_name)
num_rounds = 4

if run_method in (mcabc, smcabc, sl): # abc algorithms
kwargs = dict()
else: # neural algorithms
kwargs = dict(
num_rounds=2,
num_rounds=num_rounds,
training_batch_size=100,
neural_net="mlp" if run_method == snre else "maf",
)
Expand Down