From 013bb6877801613ca13f347b3b199b4fdd81cfd6 Mon Sep 17 00:00:00 2001 From: Imahn Date: Sat, 11 Mar 2023 12:04:05 +0100 Subject: [PATCH 1/5] For SNPE, MAF with RQS is available now. --- sbi/neural_nets/flow.py | 111 +++++++++++++++++++++++++++++++++++++ sbi/utils/get_nn_models.py | 10 +++- 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/sbi/neural_nets/flow.py b/sbi/neural_nets/flow.py index 4d2e5a31f..ae199174a 100644 --- a/sbi/neural_nets/flow.py +++ b/sbi/neural_nets/flow.py @@ -9,6 +9,7 @@ from pyknos.nflows import distributions as distributions_ from pyknos.nflows import flows, transforms from pyknos.nflows.nn import nets +from pyknos.nflows.transforms.splines import rational_quadratic from torch import Tensor, nn, relu, tanh, tensor, uint8 from sbi.utils.sbiutils import ( @@ -179,6 +180,116 @@ def build_maf( return neural_net +def build_maf_rqs( + batch_x: Tensor, + batch_y: Tensor, + z_score_x: Optional[str] = "independent", + z_score_y: Optional[str] = "independent", + hidden_features: int = 50, + num_transforms: int = 5, + embedding_net: nn.Module = nn.Identity(), + num_blocks: int = 2, + num_bins: int = 10, + tails: Optional[str] = None, + tail_bound: float = 1.0, + dropout_probability: float = 0.0, + use_batch_norm: bool = False, + min_bin_width: float = rational_quadratic.DEFAULT_MIN_BIN_WIDTH, + min_bin_height: float = rational_quadratic.DEFAULT_MIN_BIN_HEIGHT, + min_derivative: float = rational_quadratic.DEFAULT_MIN_DERIVATIVE, + **kwargs, +) -> nn.Module: + """Builds MAF p(x|y), where the diffeomorphisms are rational-quadratic + splines (RQS). + + Args: + batch_x: Batch of xs, used to infer dimensionality and (optional) z-scoring. + batch_y: Batch of ys, used to infer dimensionality and (optional) z-scoring. + z_score_x: Whether to z-score xs passing into the network, can be one of: + - `none`, or None: do not z-score. + - `independent`: z-score each dimension independently. + - `structured`: treat dimensions as related, therefore compute mean and std + over the entire batch, instead of per-dimension. Should be used when each + sample is, for example, a time series or an image. + z_score_y: Whether to z-score ys passing into the network, same options as + z_score_x. + hidden_features: Number of hidden features. + num_transforms: Number of transforms. + embedding_net: Optional embedding network for y. + num_blocks: number of blocks used for residual net for context embedding. + num_bins: Number of bins of the RQS. + tails: Whether to use constrained or unconstrained RQS, can be one of: + - None: constrained RQS. + - 'linear': unconstrained RQS (RQS transformation is only + applied on domain [-B, B], with `linear` tails, outside [-B, B], + identity transformation is returned). + tail_bound: RQS transformation is applied on domain [-B, B], + `tail_bound` is equal to B. + dropout_probability: dropout probability for regularization in residual net. + use_batch_norm: whether to use batch norm in residual net. + min_bin_width: Minimum bin width. + min_bin_height: Minimum bin height. + min_derivative: Minimum derivative at knot values of bins. + kwargs: Additional arguments that are passed by the build function but are not + relevant for maf and are therefore ignored. + + Returns: + Neural network. + """ + x_numel = batch_x[0].numel() + # Infer the output dimensionality of the embedding_net by making a forward pass. + check_data_device(batch_x, batch_y) + check_embedding_net_device(embedding_net=embedding_net, datum=batch_y) + y_numel = embedding_net(batch_y[:1]).numel() + + if x_numel == 1: + warn("In one-dimensional output space, this flow is limited to Gaussians") + + transform_list = [] + for _ in range(num_transforms): + block = [ + transforms.MaskedPiecewiseRationalQuadraticAutoregressiveTransform( + features=x_numel, + hidden_features=hidden_features, + context_features=y_numel, + num_bins=num_bins, + tails=tails, + tail_bound=tail_bound, + num_blocks=num_blocks, + use_residual_blocks=False, + random_mask=False, + activation=tanh, + dropout_probability=dropout_probability, + use_batch_norm=use_batch_norm, + min_bin_width=min_bin_width, + min_bin_height=min_bin_height, + min_derivative=min_derivative, + ), + transforms.RandomPermutation(features=x_numel), + ] + transform_list += block + + z_score_x_bool, structured_x = z_score_parser(z_score_x) + if z_score_x_bool: + transform_list = [ + standardizing_transform(batch_x, structured_x) + ] + transform_list + + z_score_y_bool, structured_y = z_score_parser(z_score_y) + if z_score_y_bool: + embedding_net = nn.Sequential( + standardizing_net(batch_y, structured_y), embedding_net + ) + + # Combine transforms. + transform = transforms.CompositeTransform(transform_list) + + distribution = distributions_.StandardNormal((x_numel,)) + neural_net = flows.Flow(transform, distribution, embedding_net) + + return neural_net + + def build_nsf( batch_x: Tensor, batch_y: Tensor, diff --git a/sbi/utils/get_nn_models.py b/sbi/utils/get_nn_models.py index 68e88ffc1..ba8f1da55 100644 --- a/sbi/utils/get_nn_models.py +++ b/sbi/utils/get_nn_models.py @@ -11,7 +11,9 @@ build_mlp_classifier, build_resnet_classifier, ) -from sbi.neural_nets.flow import build_made, build_maf, build_nsf +from sbi.neural_nets.flow import ( + build_made, build_maf, build_maf_rqs, build_nsf +) from sbi.neural_nets.mdn import build_mdn from sbi.neural_nets.mnle import build_mnle @@ -191,7 +193,7 @@ def posterior_nn( Args: model: The type of density estimator that will be created. One of [`mdn`, - `made`, `maf`, `nsf`]. + `made`, `maf`, `maf_rqs`, `nsf`]. z_score_theta: Whether to z-score parameters $\theta$ before passing them into the network, can take one of the following: - `none`, or None: do not z-score. @@ -261,6 +263,10 @@ def build_fn(batch_theta, batch_x): return build_made(batch_x=batch_theta, batch_y=batch_x, **kwargs) elif model == "maf": return build_maf(batch_x=batch_theta, batch_y=batch_x, **kwargs) + elif model == "maf_rqs": + return build_maf_rqs( + batch_x=batch_theta, batch_y=batch_x, **kwargs + ) elif model == "nsf": return build_nsf(batch_x=batch_theta, batch_y=batch_x, **kwargs) else: From c32193a74777187503f2e403327b73c5dafe724c Mon Sep 17 00:00:00 2001 From: janfb Date: Mon, 13 Mar 2023 08:58:38 +0100 Subject: [PATCH 2/5] add density estimators to tests. --- tests/inference_on_device_test.py | 1 - tests/linearGaussian_snle_test.py | 9 +++-- tests/linearGaussian_snpe_test.py | 63 +++++++++++++++++++++++++++---- 3 files changed, 61 insertions(+), 12 deletions(-) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index c18e9cd40..3acc5724e 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -24,7 +24,6 @@ SNRE_C, DirectPosterior, MCMCPosterior, - NeuralInference, RejectionPosterior, VIPosterior, likelihood_estimator_based_potential, diff --git a/tests/linearGaussian_snle_test.py b/tests/linearGaussian_snle_test.py index 1916497c9..7b687ec82 100644 --- a/tests/linearGaussian_snle_test.py +++ b/tests/linearGaussian_snle_test.py @@ -26,7 +26,8 @@ true_posterior_linear_gaussian_mvn_prior, ) from sbi.utils import likelihood_nn -from tests.test_utils import check_c2st, get_prob_outside_uniform_prior + +from .test_utils import check_c2st, get_prob_outside_uniform_prior @pytest.mark.parametrize("num_dim", (1, 3)) @@ -68,7 +69,7 @@ def test_api_snl_on_linearGaussian(num_dim: int): posterior.sample(sample_shape=(num_samples,)) -def test_c2st_snl_on_linearGaussian(): +def test_c2st_snl_on_linearGaussian(density_estimator="maf"): """Test whether SNL infers well a simple example with available ground truth. This example has different number of parameters theta than number of x. This test @@ -106,7 +107,7 @@ def test_c2st_snl_on_linearGaussian(): ), prior, ) - density_estimator = likelihood_nn("maf", num_transforms=3) + density_estimator = likelihood_nn(model=density_estimator, num_transforms=3) inference = SNLE(density_estimator=density_estimator, show_progress_bars=False) theta, x = simulate_for_sbi( @@ -127,7 +128,7 @@ def test_c2st_snl_on_linearGaussian(): samples = posterior.sample((num_samples,)) # Compute the c2st and assert it is near chance level of 0.5. - check_c2st(samples, target_samples, alg="snle_a") + check_c2st(samples, target_samples, alg=f"snle_a-{density_estimator}") @pytest.mark.slow diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py index 152c48f41..460f517c3 100644 --- a/tests/linearGaussian_snpe_test.py +++ b/tests/linearGaussian_snpe_test.py @@ -31,8 +31,9 @@ true_posterior_linear_gaussian_mvn_prior, ) from sbi.utils import RestrictedPrior, get_density_thresholder -from tests.sbiutils_test import conditional_of_mvn -from tests.test_utils import ( + +from .sbiutils_test import conditional_of_mvn +from .test_utils import ( check_c2st, get_dkl_gaussian_prior, get_normalization_uniform_prior, @@ -144,12 +145,58 @@ def test_c2st_snpe_on_linearGaussian( assert ((map_ - ones(num_dim)) ** 2).sum() < 0.5 -def test_c2st_snpe_on_linearGaussian_different_dims(): +@pytest.mark.slow +@pytest.mark.parametrize("density_estimtor", ["mdn", "maf", "nsf"]) +def test_density_estimators_on_linearGaussian(density_estimtor): + """Test SNPE with different density estimators on linear Gaussian example.""" + + theta_dim = 4 + x_dim = 4 + + x_o = zeros(1, x_dim) + num_samples = 1000 + num_simulations = 2000 + + # likelihood_mean will be likelihood_shift+theta + likelihood_shift = -1.0 * ones(x_dim) + likelihood_cov = 0.3 * eye(x_dim) + + prior_mean = zeros(theta_dim) + prior_cov = eye(theta_dim) + + prior = MultivariateNormal(loc=prior_mean, covariance_matrix=prior_cov) + 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,)) + + simulator, prior = prepare_for_sbi( + lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior + ) + + inference = SNPE_C(prior, density_estimator=density_estimtor) + + theta, x = simulate_for_sbi( + simulator, prior, num_simulations, simulation_batch_size=1000 + ) + posterior_estimator = inference.append_simulations(theta, x).train( + training_batch_size=100 + ) + posterior = DirectPosterior( + prior=prior, posterior_estimator=posterior_estimator + ).set_default_x(x_o) + samples = posterior.sample((num_samples,)) + + # Compute the c2st and assert it is near chance level of 0.5. + check_c2st(samples, target_samples, alg=f"snpe_{density_estimtor}") + + +def test_c2st_snpe_on_linearGaussian_different_dims(density_estimator="maf"): """Test whether SNPE B/C infer well a simple example with available ground truth. - This example has different number of parameters theta than number of x. Also - this implicitly tests simulation_batch_size=1. It also impleictly tests whether the - prior can be `None` and whether we can stop and resume training. + This test uses a linear Gaussian example with different number of parameters and + data dimensions. It tests different density estimators. Additionally, it implicitly + tests whether the prior can be `None` and whether we can stop and resume training. """ @@ -184,7 +231,9 @@ def test_c2st_snpe_on_linearGaussian_different_dims(): prior, ) # Test whether prior can be `None`. - inference = SNPE_C(prior=None, density_estimator="maf", show_progress_bars=False) + inference = SNPE_C( + prior=None, density_estimator=density_estimator, show_progress_bars=False + ) # type: ignore theta, x = simulate_for_sbi(simulator, prior, 2000, simulation_batch_size=1) From 6c4effda76bd272738547393006bdeb84389e1cb Mon Sep 17 00:00:00 2001 From: Imahn Date: Mon, 13 Mar 2023 21:43:26 +0100 Subject: [PATCH 3/5] For SNLE, MAF with RQS is available now; MAF with RQS added to tests; Black formatting. --- sbi/utils/get_nn_models.py | 16 +++++++--------- tests/inference_on_device_test.py | 2 ++ tests/linearGaussian_snpe_test.py | 2 +- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/sbi/utils/get_nn_models.py b/sbi/utils/get_nn_models.py index ba8f1da55..71938df37 100644 --- a/sbi/utils/get_nn_models.py +++ b/sbi/utils/get_nn_models.py @@ -11,9 +11,7 @@ build_mlp_classifier, build_resnet_classifier, ) -from sbi.neural_nets.flow import ( - build_made, build_maf, build_maf_rqs, build_nsf -) +from sbi.neural_nets.flow import build_made, build_maf, build_maf_rqs, build_nsf from sbi.neural_nets.mdn import build_mdn from sbi.neural_nets.mnle import build_mnle @@ -111,7 +109,7 @@ def likelihood_nn( Args: model: The type of density estimator that will be created. One of [`mdn`, - `made`, `maf`, `nsf`]. + `made`, `maf`, `maf_rqs`, `nsf`]. z_score_theta: Whether to z-score parameters $\theta$ before passing them into the network, can take one of the following: - `none`, or None: do not z-score. @@ -160,10 +158,12 @@ def likelihood_nn( def build_fn(batch_theta, batch_x): if model == "mdn": return build_mdn(batch_x=batch_x, batch_y=batch_theta, **kwargs) - if model == "made": + elif model == "made": return build_made(batch_x=batch_x, batch_y=batch_theta, **kwargs) - if model == "maf": + elif model == "maf": return build_maf(batch_x=batch_x, batch_y=batch_theta, **kwargs) + elif model == "maf_rqs": + return build_maf_rqs(batch_x=batch_x, batch_y=batch_theta, **kwargs) elif model == "nsf": return build_nsf(batch_x=batch_x, batch_y=batch_theta, **kwargs) elif model == "mnle": @@ -264,9 +264,7 @@ def build_fn(batch_theta, batch_x): elif model == "maf": return build_maf(batch_x=batch_theta, batch_y=batch_x, **kwargs) elif model == "maf_rqs": - return build_maf_rqs( - batch_x=batch_theta, batch_y=batch_x, **kwargs - ) + return build_maf_rqs(batch_x=batch_theta, batch_y=batch_x, **kwargs) elif model == "nsf": return build_nsf(batch_x=batch_theta, batch_y=batch_x, **kwargs) else: diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index 3acc5724e..baaa956e5 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -51,10 +51,12 @@ (SNPE_C, "mdn", "rejection"), (SNPE_C, "maf", "slice"), (SNPE_C, "maf", "direct"), + (SNPE_C, "maf_rqs", "direct"), (SNLE, "maf", "slice"), (SNLE, "nsf", "slice_np"), (SNLE, "nsf", "rejection"), (SNLE, "maf", "importance"), + (SNLE, "maf_rqs", "slice"), (SNRE_A, "mlp", "slice_np_vectorized"), (SNRE_B, "resnet", "nuts"), (SNRE_B, "resnet", "rejection"), diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py index 460f517c3..039c393a4 100644 --- a/tests/linearGaussian_snpe_test.py +++ b/tests/linearGaussian_snpe_test.py @@ -146,7 +146,7 @@ def test_c2st_snpe_on_linearGaussian( @pytest.mark.slow -@pytest.mark.parametrize("density_estimtor", ["mdn", "maf", "nsf"]) +@pytest.mark.parametrize("density_estimtor", ["mdn", "maf", "maf_rqs", "nsf"]) def test_density_estimators_on_linearGaussian(density_estimtor): """Test SNPE with different density estimators on linear Gaussian example.""" From 2a6c928e9dd7ed7cc54e564ad6b35d4c0506336f Mon Sep 17 00:00:00 2001 From: Imahn Date: Tue, 14 Mar 2023 12:08:12 +0100 Subject: [PATCH 4/5] Adjustment of the MAF RQS default values for consistency to the NSF. --- sbi/neural_nets/flow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sbi/neural_nets/flow.py b/sbi/neural_nets/flow.py index ae199174a..a3465ac8e 100644 --- a/sbi/neural_nets/flow.py +++ b/sbi/neural_nets/flow.py @@ -190,8 +190,8 @@ def build_maf_rqs( embedding_net: nn.Module = nn.Identity(), num_blocks: int = 2, num_bins: int = 10, - tails: Optional[str] = None, - tail_bound: float = 1.0, + tails: Optional[str] = 'linear', + tail_bound: float = 3.0, dropout_probability: float = 0.0, use_batch_norm: bool = False, min_bin_width: float = rational_quadratic.DEFAULT_MIN_BIN_WIDTH, From 5132a76e4a73b74a76cfe5b3d1d9f894545d8fa8 Mon Sep 17 00:00:00 2001 From: Imahn Date: Tue, 14 Mar 2023 22:42:45 +0100 Subject: [PATCH 5/5] Black formatting. --- sbi/neural_nets/flow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sbi/neural_nets/flow.py b/sbi/neural_nets/flow.py index a3465ac8e..49da8af77 100644 --- a/sbi/neural_nets/flow.py +++ b/sbi/neural_nets/flow.py @@ -190,7 +190,7 @@ def build_maf_rqs( embedding_net: nn.Module = nn.Identity(), num_blocks: int = 2, num_bins: int = 10, - tails: Optional[str] = 'linear', + tails: Optional[str] = "linear", tail_bound: float = 3.0, dropout_probability: float = 0.0, use_batch_norm: bool = False,