def test_transforms(task_name):
task = sbibm.get_task(task_name)
observation = task.get_observation(num_observation=1)
true_parameters = task.get_true_parameters(num_observation=1)
transforms = task._get_transforms(
automatic_transform_enabled=True)["parameters"]
parameters_constrained = true_parameters
parameters_unconstrained = transforms(true_parameters)
lpf_1 = task._get_log_prob_fn(observation=observation,
automatic_transform_enabled=False)
log_prob_1 = lpf_1(parameters_constrained)
lpf_2 = task._get_log_prob_fn(observation=observation,
automatic_transform_enabled=True)
# lpf_2 takes unconstrained parameters are inputs and returns
# the log prob of the unconstrained distribution
log_prob_2 = lpf_2(parameters_unconstrained)
# through change of variables, we can recover the original log prob
# ladj(x,y) -> log |dy/dx| -> ladj(untransformed, transformed)
log_prob_3 = log_prob_2 + transforms.log_abs_det_jacobian(
parameters_constrained, parameters_unconstrained)
assert torch.allclose(log_prob_1, log_prob_3)
def test_prior(
task_name,
num_observation=1,
num_samples=1000,
):
task = sbibm.get_task(task_name)
samples = run_prior(
task=task,
num_observation=num_observation,
num_samples=num_samples,
)
assert len(samples) == num_samples
def test_log_prob_fn(task_name, jit_compile, batch_size, implementation,
posterior):
"""Test `get_log_prob_fn`
Uses test cases for which the true posterior is known in closed form. Since
`log_prob_fn` returns the unnormalized posterior log probability, it is tested
whether the two are proportional.
"""
task = sbibm.get_task(task_name)
prior = task.get_prior()
prior_dist = task.get_prior_dist()
posterior_dist = task._get_reference_posterior(num_observation=1)
log_prob = task._get_log_prob_fn(
num_observation=1,
implementation=implementation,
jit_compile=jit_compile,
posterior=posterior,
)
parameters = prior(num_samples=batch_size)
# Test whether batching works
if batch_size > 1:
for b in range(batch_size):
torch.allclose(
log_prob(parameters)[b],
log_prob(parameters[b, :].reshape(1, -1)))
torch.allclose(
posterior_dist.log_prob(parameters)[b],
posterior_dist.log_prob(parameters[b, :].reshape(1, -1)),
)
# Test whether proportionality holds
diff_ref = log_prob(parameters) - posterior_dist.log_prob(parameters)
if not posterior:
diff_ref += prior_dist.log_prob(parameters)
for _ in range(10):
parameters = prior(num_samples=batch_size)
diff = log_prob(parameters) - posterior_dist.log_prob(parameters)
if not posterior:
diff += prior_dist.log_prob(parameters)
assert torch.allclose(diff, diff_ref)
def test_rejection_with_proposal(
plt,
task_name="gaussian_linear_uniform",
num_observation=1,
num_samples=10000,
prior_weight=0.1,
multiplier_M=1.2,
batch_size=10000,
num_batches_without_new_max=1000,
):
task = sbibm.get_task(task_name)
reference_samples = task.get_reference_posterior_samples(
num_observation=num_observation)
proposal_dist = get_proposal(
task=task,
samples=reference_samples,
prior_weight=prior_weight,
bounded=False,
density_estimator="flow",
flow_model="nsf",
)
samples = run(
task=task,
num_observation=num_observation,
num_samples=num_samples,
batch_size=batch_size,
num_batches_without_new_max=num_batches_without_new_max,
multiplier_M=multiplier_M,
proposal_dist=proposal_dist,
)
num_samples_plotting = 1000
pairplot([
samples.numpy()[:num_samples_plotting, :],
reference_samples.numpy()[:num_samples_plotting, :],
])
acc = c2st(samples, reference_samples[:num_samples, :])
assert torch.abs(acc - 0.5) < 0.01
def test_log_prob_grad_fn(jit_compile, batch_size, implementation):
"""Test `get_log_prob_grad_fn`
We are using the likleihood of the Gaussian linear using the fact that:
∇ wrt p of log N(p|0, 1) is -p,
since that is the derivative of -((p-0)**2)/(2.*1.).
We are checking against this analytical derivative.
"""
task = sbibm.get_task("gaussian_linear", simulator_scale=1.0)
observation = torch.zeros((10, ))
prior = task.get_prior()
log_prob_grad = task._get_log_prob_grad_fn(
observation=observation,
implementation=implementation,
jit_compile=jit_compile,
posterior=False,
)
parameters = prior(num_samples=batch_size)
# Test whether batching works
if batch_size > 1:
for b in range(batch_size):
torch.allclose(
log_prob_grad(parameters)[b],
log_prob_grad(parameters[b, :].reshape(1, -1)),
)
# Test whether gradient is correct
grads = log_prob_grad(parameters)
analytical_grad = -1.0 * parameters
assert torch.allclose(grads, analytical_grad)
def fig_posterior(
task_name: str,
num_observation: int = 1,
num_samples: int = 1000,
prior: bool = False,
reference: bool = True,
true_parameter: bool = False,
samples_path: Optional[str] = None,
samples_tensor: Optional[torch.Tensor] = None,
samples_name: Optional[str] = None,
samples_color: Optional[str] = None,
title: Optional[str] = None,
title_dx: int = 0,
legend: bool = True,
seed: int = 101,
config: Optional[str] = None,
width: Optional[int] = None,
height: Optional[int] = None,
default_color: str = "#0035FD",
colors_dict: Dict[str, Any] = {},
interactive: bool = False,
limits: Optional[Union[List[float], str]] = None,
num_bins: int = 40,
scatter_size: float = 1.0,
**kwargs: Any,
):
"""Plots posteriors samples for given task
Args:
task_name: Name of the task to plot posteriors for
num_observation: Observation number
num_samples: Number of samples to use for plotting
prior: Whether or not to plot prior samples
reference: Whether or not to plot reference posterior samples
samples_path: If specified, will load samples from disk from path
samples_tensor: Instead of a path, samples can also be passed as torch.Tensor
samples_name: Name for samples, defaults to "Algorithm"
samples_color: Optional string for color of samples
title: Title for plot
title_dx: x-direction offset for title
legend: Whether to plot a legend
seed: Seed
width: Width
height: Height
default_color: Default color of samples
colors_dict: Dictionary of colors
config: Optional string to load predefined config
interactive: Interactive mode (experimental)
limits: Optional limits, can also be a string, e.g., "Prior", "Ref. Posterior",
"Algorithm" / `samples_name`, in which case respective ranges are used
num_bins: Number of bins
scatter_size: Scatter size
Returns:
Chart
"""
# Samples to plot
task = sbibm.get_task(task_name)
samples = []
labels_samples = []
colors = {}
samples_prior = task.get_prior()(num_samples=num_samples)
if prior:
sample_name = "Prior"
samples.append(samples_prior.numpy())
labels_samples.append(sample_name)
if sample_name in colors_dict:
colors[sample_name] = colors_dict[sample_name]
else:
colors[sample_name] = "#646464"
if reference:
sample_name = "Ref. Posterior"
samples_reference = sample(
task.get_reference_posterior_samples(
num_observation=num_observation).numpy(),
num_samples,
replace=False,
seed=seed,
)
samples.append(samples_reference)
labels_samples.append(sample_name)
if sample_name in colors_dict:
colors[sample_name] = colors_dict[sample_name]
else:
colors[sample_name] = "#0a0a0a"
if true_parameter:
sample_name = "True parameter"
samples.append(
task.get_true_parameters(num_observation=num_observation).repeat(
num_samples, 1).numpy())
labels_samples.append(sample_name)
if sample_name in colors_dict:
colors[sample_name] = colors_dict[sample_name]
else:
colors[sample_name] = "#f92700"
if samples_tensor is not None or samples_path is not None:
if samples_tensor is not None:
samples_ = samples_tensor.numpy()
else:
samples_ = get_ndarray_from_csv(samples_path)
samples_algorithm = sample(samples_,
num_samples,
replace=False,
seed=seed)
samples.append(samples_algorithm)
if samples_name is None:
sample_name = "Algorithm"
else:
sample_name = samples_name
labels_samples.append(sample_name)
if samples_color is not None:
colors[sample_name] = samples_color
else:
if sample_name in colors_dict:
colors[sample_name] = colors_dict[samples_name]
else:
colors[sample_name] = default_color
if len(samples) == 0:
return None
for s in samples:
assert s.shape[0] == num_samples
numbers_unicode = ["₀", "₁", "₂", "₃", "₄", "₅", "₆", "₇", "₈", "₉", "₁₀"]
labels_dim = [
f"θ{numbers_unicode[i+1]}" for i in range(task.dim_parameters)
]
df = den.np2df(
samples=[sample for sample in samples],
field="sample",
labels_samples=labels_samples,
labels_dim=labels_dim,
)
style = {}
keywords = {}
keywords["color"] = den.colorscale(colors,
shorthand="sample:N",
legend=legend)
keywords["interactive"] = interactive
if limits is None:
if task_name in _LIMITS_:
limits = _LIMITS_[task_name]
else:
limits = [
list(i) for i in zip(
samples_prior.min(dim=0)[0].tolist(),
samples_prior.max(dim=0)[0].tolist(),
)
]
elif type(limits) == str:
assert limits in labels_samples
samples_limits = torch.from_numpy(
samples[labels_samples.index(limits)])
limits = [
list(i) for i in zip(
samples_limits.min(dim=0)[0].tolist(),
samples_limits.max(dim=0)[0].tolist(),
)
]
keywords["limits"] = limits
keywords["num_bins"] = num_bins
if config == "manuscript":
style["font_family"] = "Inter"
keywords["width"] = 100 if width is None else width
keywords["height"] = 100 if height is None else height
style["font_size"] = 12
if config == "streamlit":
size = 500 / task.dim_parameters
keywords["width"] = size if width is None else width
keywords["height"] = size if height is None else height
style["font_size"] = 16
alt.themes.enable("default")
den.set_style(
extra={
"config": {
"axisX": {
"domain": False,
"domainWidth": 0,
"ticks": False,
"tickWidth": 0,
"grid": False,
},
"axisY": {
"domain": False,
"domainWidth": 0,
"ticks": False,
"tickWidth": 0,
"grid": False,
},
}
},
**style,
)
chart = den.pairplot(
df,
field="sample",
scatter_size=scatter_size,
bar_opacity=0.4,
**keywords,
)
if title is not None:
chart = chart.properties(title={
"text": [title],
}).configure_title(offset=10,
orient="top",
anchor="middle",
dx=title_dx)
return chart
import sbibm
from sbibm.algorithms.pytorch.baseline_sir import run
from sbibm.algorithms.pytorch.utils.proposal import get_proposal
from sbibm.metrics.c2st import c2st
def test_sir_with_proposal(
plt,
task_name="gaussian_linear_uniform",
num_observation=1,
num_samples=10000,
num_simulations=10_000_000,
batch_size=10000,
prior_weight=0.01,
):
task = sbibm.get_task(task_name)
reference_samples = task.get_reference_posterior_samples(
num_observation=num_observation)
proposal_dist = get_proposal(
task=task,
samples=reference_samples,
prior_weight=prior_weight,
bounded=False,
density_estimator="flow",
flow_model="nsf",
)
samples = run(
task=task,
def main(cfg: DictConfig) -> None:
log = logging.getLogger(__name__)
log.info(cfg.pretty())
log.info(f"sbibm version: {sbibm.__version__}")
log.info(f"Hostname: {socket.gethostname()}")
if cfg.seed is None:
log.info(
"Seed not specified, generating random seed for replicability")
cfg.seed = int(torch.randint(low=1, high=2**32 - 1, size=(1, ))[0])
log.info(f"Random seed: {cfg.seed}")
save_config(cfg)
# Seeding
torch.manual_seed(cfg.seed)
random.seed(cfg.seed)
np.random.seed(cfg.seed)
# Devices
gpu = True if cfg.device != "cpu" else False
if gpu:
torch.cuda.set_device(0)
torch.set_default_tensor_type(
"torch.cuda.FloatTensor" if gpu else "torch.FloatTensor")
# Paths
path_samples = "posterior_samples.csv.bz2"
path_runtime = "runtime.csv"
path_log_prob_true_parameters = "log_prob_true_parameters.csv"
path_num_simulations_simulator = "num_simulations_simulator.csv"
path_predictive_samples = "predictive_samples.csv.bz2"
# Run
task = sbibm.get_task(cfg.task.name)
t0 = time.time()
parts = cfg.algorithm.run.split(".")
module_name = ".".join(["sbibm", "algorithms"] + parts[:-1])
run_fn = getattr(importlib.import_module(module_name), parts[-1])
algorithm_params = cfg.algorithm.params if "params" in cfg.algorithm else {}
log.info("Start run")
outputs = run_fn(
task,
num_observation=cfg.task.num_observation,
num_samples=task.num_posterior_samples,
num_simulations=cfg.task.num_simulations,
**algorithm_params,
)
runtime = time.time() - t0
log.info("Finished run")
# Store outputs
if type(outputs) == torch.Tensor:
samples = outputs
num_simulations_simulator = float("nan")
log_prob_true_parameters = float("nan")
elif type(outputs) == tuple and len(outputs) == 3:
samples = outputs[0]
num_simulations_simulator = float(outputs[1])
log_prob_true_parameters = (float(outputs[2]) if outputs[2] is not None
else float("nan"))
else:
raise NotImplementedError
save_tensor_to_csv(path_samples,
samples,
columns=task.get_labels_parameters())
save_float_to_csv(path_runtime, runtime)
save_float_to_csv(path_num_simulations_simulator,
num_simulations_simulator)
save_float_to_csv(path_log_prob_true_parameters, log_prob_true_parameters)
# Predictive samples
log.info("Draw posterior predictive samples")
simulator = task.get_simulator()
predictive_samples = []
batch_size = 1_000
for idx in range(int(samples.shape[0] / batch_size)):
try:
predictive_samples.append(
simulator(samples[(idx * batch_size):((idx + 1) *
batch_size), :]))
except:
predictive_samples.append(
float("nan") * torch.ones((batch_size, task.dim_data)))
predictive_samples = torch.cat(predictive_samples, dim=0)
save_tensor_to_csv(path_predictive_samples, predictive_samples,
task.get_labels_data())
# Compute metrics
if cfg.compute_metrics:
df_metrics = compute_metrics_df(
task_name=cfg.task.name,
num_observation=cfg.task.num_observation,
path_samples=path_samples,
path_runtime=path_runtime,
path_predictive_samples=path_predictive_samples,
path_log_prob_true_parameters=path_log_prob_true_parameters,
log=log,
)
df_metrics.to_csv("metrics.csv", index=False)
log.info(f"Metrics:\n{df_metrics.transpose().to_string(header=False)}")
def compute_metrics_df(
task_name: str,
num_observation: int,
path_samples: str,
path_runtime: str,
path_predictive_samples: str,
path_log_prob_true_parameters: str,
log: logging.Logger = logging.getLogger(__name__),
) -> pd.DataFrame:
"""Compute all metrics, returns dataframe
Args:
task_name: Task
num_observation: Observation
path_samples: Path to posterior samples
path_runtime: Path to runtime file
path_predictive_samples: Path to predictive samples
path_log_prob_true_parameters: Path to NLTP
log: Logger
Returns:
Dataframe with results
"""
log.info(f"Compute all metrics")
# Load task
task = sbibm.get_task(task_name)
# Load samples
reference_posterior_samples = task.get_reference_posterior_samples(
num_observation)[:task.num_posterior_samples, :]
algorithm_posterior_samples = get_tensor_from_csv(
path_samples)[:task.num_posterior_samples, :]
assert reference_posterior_samples.shape[0] == task.num_posterior_samples
assert algorithm_posterior_samples.shape[0] == task.num_posterior_samples
log.info(
f"Loaded {task.num_posterior_samples} samples from reference and algorithm"
)
# Load posterior predictive samples
predictive_samples = get_tensor_from_csv(
path_predictive_samples)[:task.num_posterior_samples, :]
assert predictive_samples.shape[0] == task.num_posterior_samples
# Load observation
observation = task.get_observation(num_observation=num_observation) # noqa
# Get runtime info
runtime_sec = torch.tensor(get_float_from_csv(path_runtime)) # noqa
# Get log prob true parameters
log_prob_true_parameters = torch.tensor(
get_float_from_csv(path_log_prob_true_parameters)) # noqa
# Names of all metrics as keys, values are calls that are passed to eval
# NOTE: Originally, we computed a large number of metrics, as reflected in the
# dictionary below. Ultimately, we used 10k samples and z-scoring for C2ST but
# not for MMD. If you were to adapt this code for your own pipeline of experiments,
# the entries for C2ST_Z, MMD and RT would probably suffice (and save compute).
_METRICS_ = {
#
# 10k samples
#
"C2ST":
"metrics.c2st(X=reference_posterior_samples, Y=algorithm_posterior_samples, z_score=False)",
"C2ST_Z":
"metrics.c2st(X=reference_posterior_samples, Y=algorithm_posterior_samples, z_score=True)",
"MMD":
"metrics.mmd(X=reference_posterior_samples, Y=algorithm_posterior_samples, z_score=False)",
"MMD_Z":
"metrics.mmd(X=reference_posterior_samples, Y=algorithm_posterior_samples, z_score=True)",
"KSD_GAUSS":
"metrics.ksd(task=task, num_observation=num_observation, samples=algorithm_posterior_samples, sig2=float(torch.median(torch.pdist(reference_posterior_samples))**2), log=False)",
"MEDDIST": "metrics.median_distance(predictive_samples, observation)",
#
# 1K samples
#
"C2ST_1K":
"metrics.c2st(X=reference_posterior_samples[:1000,:], Y=algorithm_posterior_samples[:1000,:], z_score=False)",
"C2ST_1K_Z":
"metrics.c2st(X=reference_posterior_samples[:1000,:], Y=algorithm_posterior_samples[:1000, :], z_score=True)",
"MMD_1K":
"metrics.mmd(X=reference_posterior_samples[:1000,:], Y=algorithm_posterior_samples[:1000, :], z_score=False)",
"MMD_1K_Z":
"metrics.mmd(X=reference_posterior_samples[:1000,:], Y=algorithm_posterior_samples[:1000, :], z_score=True)",
"KSD_GAUSS_1K":
"metrics.ksd(task=task, num_observation=num_observation, samples=algorithm_posterior_samples[:1000, :], sig2=float(torch.median(torch.pdist(reference_posterior_samples))**2), log=False)",
"MEDDIST_1K":
"metrics.median_distance(predictive_samples[:1000,:], observation)",
#
# Not based on samples
#
"NLTP": "-1. * log_prob_true_parameters",
"RT": "runtime_sec",
}
import sbibm.metrics as metrics # noqa
metrics_dict = {}
for metric, eval_cmd in _METRICS_.items():
log.info(f"Computing {metric}")
try:
metrics_dict[metric] = eval(eval_cmd).cpu().numpy().astype(
np.float32)
log.info(f"{metric}: {metrics_dict[metric]}")
except:
metrics_dict[metric] = float("nan")
return pd.DataFrame(metrics_dict)