lf2i.test_statistics package¶

Note

TestStatistic, ACORE, BFF, Waldo and Posterior are re-exported at the lf2i.test_statistics package level; see their definitions in the submodules below.

Submodules¶

lf2i.test_statistics.acore module¶

class lf2i.test_statistics.acore.ACORE(estimator: str | Any, poi_dim: int, nuisance_dim: int, batch_size: int, data_dim: int, estimator_kwargs: Dict = {}, verbose: bool = True, n_jobs: int = -2, optimizer: str = 'coordinate_descent', param_space_bounds: List[Tuple[float]] = None, max_iter: int | None = 1, estimator_train_kwargs: Dict | None = None)[source]¶

Bases: TestStatistic

Implements the ACORE test statistic as described in https://proceedings.mlr.press/v119/dalmasso20a.html and https://arxiv.org/abs/2107.03920.

Parameters:

estimator (Union[str, Any]) – Probabilistic classifier used to estimate odds (i.e., likelihood up to a normalization constant). If str, must be one of the predefined estimators listed in test_statistics/estimators.py. If Any, a trained estimator is expected. Needs to implement estimator.predict_proba(X=…).
poi_dim (int) – Dimensionality (number) of the parameters of interest.
nuisance_dim (int) – Dimensionality (number) of the nuisance parameters (systematics). Should be 0 if all parameters are object of inference.
batch_size (int) – Size of a batch of datapoints from a specific parameter configuration. Must be the same for observations and simulations. A simulated/observed batch from a specific parameter configuration will have dimensions (batch_size, data_dim).
data_dim (int) – Dimensionality of a single datapoint X.
estimator_kwargs (Dict, optional) – Hyperparameters and settings for the conditional mean estimator, by default {}.
verbose (bool, optional) – Whether to print progress bars when evaluating or not, by default True.
n_jobs (int, optional) – Number of workers to use when computing ACORE over multiple inputs, by default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.
optimizer (str, optional) – Choice between ‘coordinate_descent’ and ‘Nelder-Mead’, by default ‘coordinate_descent’
param_space_bounds (List, optional) – List of tuples bounding optimization in each parameter component
max_iter (int, optional) – Number of iterations of coordinate descent
estimator_train_kwargs (Dict, optional) – Keyword arguments to be passed to the estimator constructor at estimation stage

estimate(parameters: ndarray | Tensor, samples: ndarray | Tensor) → None[source]¶

Train the estimator for odds (i.e. likelihood up to a normalization constant). The training dataset should contain two classes:

label 1, with pairs \((\theta, X)\) where \(X \sim p(\cdot;\theta)\) is drawn from the likelihood/simulator.
label 0, with pairs \((\theta, X)\) where \(X \sim G\) is drawn from a dominating reference distribution (e.g., empirical marginal).

To goal is to train a classifier that is able to distinguish whether a sample comes from the likelihood or not. See https://arxiv.org/abs/2107.03920 for a more detailed explanation.

Parameters:

parameters (Union[np.ndarray, torch.Tensor]) – Simulated parameters to be used for training.
samples (Union[np.ndarray, torch.Tensor]) – Simulated samples to be used for training.

evaluate(parameters: ndarray | Tensor, samples: ndarray | Tensor, mode: str, param_space_bounds: List[Tuple[float]] | None = None) → ndarray[source]¶

Evaluate the ACORE test statistic over the given parameters and samples. Behaviour differs depending on mode:

‘critical_values’ and ‘diagnostics’ compute ACORE once for each pair \((\theta, X)\).
‘confidence_sets’ computes ACORE over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (Union[np.ndarray, torch.Tensor]) – Parameters over which to evaluate the test statistic.
samples (Union[np.ndarray, torch.Tensor]) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.
param_space_bounds (Optional[List[Tuple[float]]]) – Bounds of the parameter space, both POIs and nuisances. Must be in the same order as in parameters.

Returns:

np.ndarray – ACORE test statistics evaluated over parameters and samples.

Raises:

ValueError – If mode is not among the pre-specified values.

lf2i.test_statistics.bff module¶

class lf2i.test_statistics.bff.BFF(estimator: str | Any, poi_dim: int, nuisance_dim: int, batch_size: int, data_dim: int, estimator_kwargs: Dict = {}, verbose: bool = True, n_jobs: int = -2, param_space_bounds: List[Tuple[float]] | None = None)[source]¶

Bases: TestStatistic

Implements the BFF test statistic as described in https://arxiv.org/abs/2107.03920. NOTE: for now supports only box uniform proposal distributions over the parameter space.

Parameters:

estimator (Union[str, Any]) – Probabilistic classifier used to estimate odds (i.e., likelihood up to a normalization constant). If str, must be one of the predefined estimators listed in test_statistics/_estimators.py. If Any, a trained estimator is expected. Needs to implement estimator.predict_proba(X=…).
poi_dim (int) – Dimensionality (number) of the parameters of interest.
nuisance_dim (int) – Dimensionality (number) of the nuisance parameters (systematics). Should be 0 if all parameters are object of inference.
batch_size (int) – Size of a batch of datapoints from a specific parameter configuration. Must be the same for observations and simulations. A simulated/observed batch from a specific parameter configuration will have dimensions (batch_size, data_dim).
data_dim (int) – Dimensionality of a single datapoint X.
estimator_kwargs (Dict, optional) – Hyperparameters and settings for the conditional mean estimator, by default {}.
verbose (bool, optional) – Whether to print progress bars when evaluating or not, by default True.
n_jobs (int, optional) – Number of workers to use when computing BFF over multiple inputs, by default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.
param_space_bounds (List[Tuple[float]], optional) – Bounds of the parameter space (POIs and nuisances), used as the fallback whenever evaluate(…) is not given its own param_space_bounds (e.g. when called through LF2I.inference(…)).

estimate(parameters: ndarray | Tensor, samples: ndarray | Tensor) → None[source]¶

Train the estimator for odds (i.e. likelihood up to a normalization constant).

The training dataset is created by:

label 1: pairs \((\theta, X)\) where \(X \sim p(\cdot;\theta)\) from the true joint distribution (original matched pairs).
label 0: pairs \((\theta', X)\) where \(\theta'\) is a permuted parameter vector, ensuring no overlap with the positive class pairs.

This creates a classifier that distinguishes true parameter-sample pairs from mismatched pairs, effectively learning the likelihood ratio.

See https://arxiv.org/abs/2107.03920 for a more detailed explanation.

Parameters:

parameters (Union[np.ndarray, torch.Tensor]) – Simulated parameters from the true joint distribution (n_samples, param_dim).
samples (Union[np.ndarray, torch.Tensor]) – Simulated samples from the true joint distribution (n_samples, sample_dim).

evaluate(parameters: ndarray | Tensor, samples: ndarray | Tensor, mode: str, param_space_bounds: List[Tuple[float]] | None = None) → ndarray[source]¶

Evaluate the BFF test statistic over the given parameters and samples. Behaviour differs depending on mode:

‘critical_values’ and ‘diagnostics’ compute BFF once for each pair \((\theta, X)\).
‘confidence_sets’ computes BFF over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (Union[np.ndarray, torch.Tensor]) – Parameters over which to evaluate the test statistic.
samples (Union[np.ndarray, torch.Tensor]) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.
param_space_bounds (Optional[List[Tuple[float]]]) – Bounds of the parameter space, both POIs and nuisances. Must be in the same order as in parameters. NOTE: Bounds are needed because we support only box uniform proposal distributions over the parameter space at the moment.

Returns:

np.ndarray – BFF test statistics evaluated over parameters and samples.

Raises:

ValueError – If mode is not among the pre-specified values.

lf2i.test_statistics.posterior module¶

class lf2i.test_statistics.posterior.Posterior(poi_dim: int, estimator: str | Any, estimator_kwargs: Dict = {}, n_jobs: int = -2, **posterior_kwargs)[source]¶

Bases: TestStatistic

Implements the Posterior test statistic, i.e. the (log) posterior density \(\log p(\theta \mid x)\) evaluated at a neural posterior estimator, as described in https://doi.org/10.1088/2632-2153/ae67cd.

Unlike Waldo’s posterior mode (which reduces the posterior to a conditional mean and variance via Monte Carlo sampling), this test statistic uses the posterior density directly: large values indicate that \(\theta\) is well-supported by the estimated posterior given x, so the acceptance region is right.

Parameters:

poi_dim (int) – Dimensionality (number) of the parameters of interest.
estimator (Union[str, Any]) – Neural posterior estimator. Currently compatible with posterior objects implementing the interface of lf2i.estimators.base_posteriors.AbstractNeuralPosteriorTrainer (e.g., estimators from the sbi library). If str, must be one of the predefined estimators listed in test_statistics/_estimators.py.
estimator_kwargs (Dict, optional) – Hyperparameters and settings for the posterior estimator, by default {}.
n_jobs (int, optional) – Number of workers to use when evaluating the test statistic over multiple inputs, by default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.
**posterior_kwargs (Any) – Additional keyword arguments forwarded to estimator.log_prob(theta=…, x=…, **posterior_kwargs) at evaluation time.

estimate(parameters: Tensor, samples: Tensor) → None[source]¶

Train the neural posterior estimator.

Parameters:

parameters (torch.Tensor) – Simulated parameters to be used for training.
samples (torch.Tensor) – Simulated samples to be used for training.

evaluate(parameters: Tensor, samples: Tensor, mode: str) → ndarray[source]¶

Evaluate the Posterior test statistic, i.e. \(\log p(\theta \mid x)\), over the given parameters and samples.

Behaviour differs depending on mode:

‘critical_values’ and ‘diagnostics’ evaluate the log-posterior once for each pair \((\theta, x)\).
‘confidence_sets’ evaluates the log-posterior over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (torch.Tensor) – Parameters over which to evaluate the test statistic.
samples (torch.Tensor) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.

Returns:

np.ndarray – Log-posterior density evaluated over parameters and samples.

Raises:

ValueError – If mode is not among the pre-specified values.

class lf2i.test_statistics.posterior.PosteriorPriorRatio(poi_dim: int, prior: Distribution | Any, estimator: str | Any, estimator_kwargs: Dict = {}, n_jobs: int = -2, **posterior_kwargs)[source]¶

Bases: TestStatistic

Implements the PosteriorPriorRatio test statistic, i.e. the log-ratio of the estimated posterior density to the prior density, \(\log \frac{p(\theta \mid x)}{p(\theta)}\), evaluated at a neural posterior estimator.

Large values indicate that the posterior places much more mass on \(\theta\) than the prior did, i.e. that the data \(x\) is highly informative about \(\theta\) relative to the prior; equivalently, that the prior (denominator) is very low relative to the posterior (numerator). The acceptance region is right.

Parameters:

poi_dim (int) – Dimensionality (number) of the parameters of interest.
prior (Union[torch.distributions.Distribution, Any]) – Prior distribution over the parameters of interest. Must implement log_prob(theta).
estimator (Union[str, Any]) – Neural posterior estimator. Currently compatible with posterior objects implementing the interface of lf2i.estimators.base_posteriors.AbstractNeuralPosteriorTrainer (e.g., estimators from the sbi library). If str, must be one of the predefined estimators listed in test_statistics/_estimators.py.
estimator_kwargs (Dict, optional) – Hyperparameters and settings for the posterior estimator, by default {}.
n_jobs (int, optional) – Number of workers to use when evaluating the test statistic over multiple inputs, by default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.
**posterior_kwargs (Any) – Additional keyword arguments forwarded to estimator.log_prob(theta=…, x=…, **posterior_kwargs) at evaluation time.

estimate(parameters: Tensor, samples: Tensor) → None[source]¶

Train the neural posterior estimator.

Parameters:

parameters (torch.Tensor) – Simulated parameters to be used for training.
samples (torch.Tensor) – Simulated samples to be used for training.

evaluate(parameters: Tensor, samples: Tensor, mode: str) → ndarray[source]¶

Evaluate the PosteriorPriorRatio test statistic, i.e. \(\log \frac{p(\theta \mid x)}{p(\theta)}\), over the given parameters and samples.

Behaviour differs depending on mode:

‘critical_values’ and ‘diagnostics’ evaluate the log-ratio once for each pair \((\theta, x)\).
‘confidence_sets’ evaluates the log-ratio over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (torch.Tensor) – Parameters over which to evaluate the test statistic.
samples (torch.Tensor) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.

Returns:

np.ndarray – Log posterior-to-prior ratio evaluated over parameters and samples.

Raises:

ValueError – If mode is not among the pre-specified values.

class lf2i.test_statistics.posterior.PriorPosteriorRatio(poi_dim: int, prior: Distribution | Any, estimator: str | Any, estimator_kwargs: Dict = {}, n_jobs: int = -2, **posterior_kwargs)[source]¶

Bases: TestStatistic

Implements the PriorPosteriorRatio test statistic, i.e. the log-ratio of the prior density to the estimated posterior density, \(\log \frac{p(\theta)}{p(\theta \mid x)}\), evaluated at a neural posterior estimator.

This is the negative of PosteriorPriorRatio’s statistic (\(\log \frac{p(\theta)}{p(\theta \mid x)} = -\log \frac{p(\theta \mid x)}{p(\theta)}\)), provided for workflows that require an acceptance region on the left (small values accepted, i.e. when the posterior is high relative to the prior, or equivalently when the prior is low relative to the posterior) rather than on the right, e.g. for consistency with Waldo, which also uses acceptance_region=’left’.

Parameters:

poi_dim (int) – Dimensionality (number) of the parameters of interest.
prior (Union[torch.distributions.Distribution, Any]) – Prior distribution over the parameters of interest. Must implement log_prob(theta).
estimator (Union[str, Any]) – Neural posterior estimator. Currently compatible with posterior objects implementing the interface of lf2i.estimators.base_posteriors.AbstractNeuralPosteriorTrainer (e.g., estimators from the sbi library). If str, must be one of the predefined estimators listed in test_statistics/_estimators.py.
estimator_kwargs (Dict, optional) – Hyperparameters and settings for the posterior estimator, by default {}.
n_jobs (int, optional) – Number of workers to use when evaluating the test statistic over multiple inputs, by default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.
**posterior_kwargs (Any) – Additional keyword arguments forwarded to estimator.log_prob(theta=…, x=…, **posterior_kwargs) at evaluation time.

estimate(parameters: Tensor, samples: Tensor) → None[source]¶

Train the neural posterior estimator.

Parameters:

parameters (torch.Tensor) – Simulated parameters to be used for training.
samples (torch.Tensor) – Simulated samples to be used for training.

evaluate(parameters: Tensor, samples: Tensor, mode: str) → ndarray[source]¶

Evaluate the PriorPosteriorRatio test statistic, i.e. \(\log \frac{p(\theta)}{p(\theta \mid x)}\), over the given parameters and samples.

Behaviour differs depending on mode:

‘critical_values’ and ‘diagnostics’ evaluate the log-ratio once for each pair \((\theta, x)\).
‘confidence_sets’ evaluates the log-ratio over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (torch.Tensor) – Parameters over which to evaluate the test statistic.
samples (torch.Tensor) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.

Returns:

np.ndarray – Log prior-to-posterior ratio evaluated over parameters and samples.

Raises:

ValueError – If mode is not among the pre-specified values.

lf2i.test_statistics.waldo module¶

class lf2i.test_statistics.waldo.Waldo(estimator: str | Any, poi_dim: int, estimation_method: str, num_posterior_samples: int | None = None, cond_variance_estimator: str | Any | None = None, estimator_kwargs: Dict = {}, cond_variance_estimator_kwargs: Dict = {}, verbose: bool = True, n_jobs: int = -2)[source]¶

Bases: TestStatistic

Implements the Waldo test statistic, as described in arXiv:2205.15680.

Parameters:

estimator (Union[str, Any]) –
If estimation_method == prediction, then this is the conditional mean estimator. If estimation_method == posterior, then this is the posterior estimator. Currently compatible with posterior objects with interface of lf2i.estimators.base_posteriors.AbstractNeuralPosteriorTrainer.

If str, will use one of the predefined estimators. If Any, a trained estimator is expected. Needs to implement estimator.predict(X=…) (“prediction”), or estimator.sample(sample_shape=…, x=…) (“posterior”).
poi_dim (int) – Dimensionality (number) of the parameters of interest.
estimation_method (str) – Whether the estimator is a prediction algorithm (“prediction”) or a posterior estimator (“posterior”).
num_posterior_samples (Optional[int], optional) – Number of posterior samples to draw to approximate conditional mean and variance if estimation_method == posterior, by default None
cond_variance_estimator (Optional[Union[str, Any]], optional) – If estimation_method == prediction, then this is the conditional variance estimator, by default None
estimator_kwargs (Dict) – Hyperparameters and settings for the conditional mean estimator, by default {}.
cond_variance_estimator_kwargs (Dict) – Hyperparameters and settings for the conditional variance estimator, by default {}.
verbose (bool, optional) – Whether to print progress bars when evaluating or not, by default True.
n_jobs (int, optional) – Number of workers to use when evaluating Waldo over multiple inputs if using a posterior estimator. By default -2, which uses all cores minus one. n_jobs == -1 uses all cores. If n_jobs < -1, then n_jobs = os.cpu_count()+1+n_jobs.

estimate(parameters: ndarray | Tensor, samples: ndarray | Tensor) → None[source]¶

Train the estimator(s) for the conditional mean and conditional variance.

Parameters:

parameters (Union[np.ndarray, torch.Tensor]) – Simulated parameters to be used for training.
samples (Union[np.ndarray, torch.Tensor]) – Simulated samples to be used for training.

evaluate(parameters: ndarray | Tensor, samples: ndarray | Tensor, mode: str) → ndarray[source]¶

Evaluate the Waldo test statistic over the given parameters and samples.

Behaviour differs depending on mode: ‘critical_values’, ‘confidence_sets’, ‘diagnostics’:

If mode equals critical_values or diagnostics, evaluate Waldo over pairs \((\theta_i, x_i)\).
If mode equals confidence_sets, evaluate Waldo over all pairs given by the cartesian product of parameters (the parameter grid to construct confidence sets) and samples.

Parameters:

parameters (np.ndarray) – Parameters over which to evaluate the test statistic.
samples (np.ndarray) – Samples over which to evaluate the test statistic.
mode (str) – Either ‘critical_values’, ‘confidence_sets’, ‘diagnostics’.

Returns:

np.ndarray – Waldo test statistics evaluated over parameters and samples.