Movatterモバイル変換


[0]ホーム

URL:


Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Refactorsimulate_for_sbi location#1253

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.

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

Already on GitHub?Sign in to your account

Merged
janfb merged 4 commits intosbi-dev:mainfromsamadpls:refactor/simulate_for_sbi-location
Sep 3, 2024
Merged
Changes from1 commit
Commits
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
PrevPrevious commit
NextNext commit
Added to utilis, Fixed Typos
Signed-off-by: samadpls <abdulsamadsid1@gmail.com>
  • Loading branch information
@samadpls
samadpls committedSep 2, 2024
commitfc1ecdfc1b0ec3a5ffee51b38f7e73808b732915
119 changes: 119 additions & 0 deletionssbi/utils/simulation_utils.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>

from typing import Any, Callable, Optional, Tuple, Union

import numpy as np
import torch
from joblib import Parallel, delayed
from numpy import ndarray
from torch import Tensor, float32
from tqdm.auto import tqdm

from sbi.utils.sbiutils import seed_all_backends


# Refactoring following #1175. tl:dr: letting joblib iterate over numpy arrays
# allows for a roughly 10x performance gain. The resulting casting necessity
# (cfr. user_input_checks.wrap_as_joblib_efficient_simulator) introduces
# considerable overhead. The simulating pipeline should, therefore, be further
# restructured in the future (PR #1188).
def simulate_for_sbi(
simulator: Callable,
proposal: Any,
num_simulations: int,
num_workers: int = 1,
simulation_batch_size: Union[int, None] = 1,
seed: Optional[int] = None,
show_progress_bar: bool = True,
) -> Tuple[Tensor, Tensor]:
r"""Returns ($\theta, x$) pairs obtained from sampling the proposal and simulating.

This function performs two steps:

- Sample parameters $\theta$ from the `proposal`.
- Simulate these parameters to obtain $x$.

Args:
simulator: A function that takes parameters $\theta$ and maps them to
simulations, or observations, `x`, $\text{sim}(\theta)\to x$. Any
regular Python callable (i.e. function or class with `__call__` method)
can be used. Note that the simulator should be able to handle numpy
arrays for efficient parallelization. You can use
`process_simulator` to ensure this.
proposal: Probability distribution that the parameters $\theta$ are sampled
from.
num_simulations: Number of simulations that are run.
num_workers: Number of parallel workers to use for simulations.
simulation_batch_size: Number of parameter sets of shape
(simulation_batch_size, parameter_dimension) that the simulator
receives per call. If None, we set
simulation_batch_size=num_simulations and simulate all parameter
sets with one call. Otherwise, we construct batches of parameter
sets and distribute them among num_workers.
seed: Seed for reproducibility.
show_progress_bar: Whether to show a progress bar for simulating. This will not
affect whether there will be a progressbar while drawing samples from the
proposal.

Returns: Sampled parameters $\theta$ and simulation-outputs $x$.
"""

if num_simulations == 0:
theta = torch.tensor([], dtype=float32)
x = torch.tensor([], dtype=float32)

else:
# Cast theta to numpy for better joblib performance (seee #1175)
seed_all_backends(seed)
theta = proposal.sample((num_simulations,))

# Parse the simulation_batch_size logic
if simulation_batch_size is None:
simulation_batch_size = num_simulations
else:
simulation_batch_size = min(simulation_batch_size, num_simulations)

if num_workers != 1:
# For multiprocessing, we want to switch to numpy arrays.
# The batch size will be an approximation, since np.array_split does
# not take as argument the size of the batch but their total.
num_batches = num_simulations // simulation_batch_size
batches = np.array_split(theta.numpy(), num_batches, axis=0)
batch_seeds = np.random.randint(low=0, high=1_000_000, size=(len(batches),))

# define seeded simulator.
def simulator_seeded(theta: ndarray, seed: int) -> Tensor:
seed_all_backends(seed)
return simulator(theta)

try: # catch TypeError to give more informative error message
simulation_outputs: list[Tensor] = [ # pyright: ignore
xx
for xx in tqdm(
Parallel(return_as="generator", n_jobs=num_workers)(
delayed(simulator_seeded)(batch, seed)
for batch, seed in zip(batches, batch_seeds)
),
total=num_simulations,
disable=not show_progress_bar,
)
]
except TypeError as err:
raise TypeError(
"For multiprocessing, we switch to numpy arrays. Make sure to "
"preprocess your simulator with `process_simulator` to handle numpy"
" arrays."
) from err

else:
simulation_outputs: list[Tensor] = []
batches = torch.split(theta, simulation_batch_size)
for batch in tqdm(batches, disable=not show_progress_bar):
simulation_outputs.append(simulator(batch))

# Correctly format the output
x = torch.cat(simulation_outputs, dim=0)
theta = torch.as_tensor(theta, dtype=float32)

return theta, x

[8]ページ先頭

©2009-2025 Movatter.jp