API reference

Bayesian calibration of simulations using Gaussian process emulation.

acquisition_functions

Acquisition functions for selecting new inputs points to evaluate model at.

get_expected_integrated_variance_acquisition_function

get_expected_integrated_variance_acquisition_function(
    gp_mean_and_variance,
    gp_lookahead_variance_reduction,
    integration_inputs,
    integration_log_weights,
)

Construct acquisition function for minimising expected integrated variance.

Selects next input points to evaluate log-likelihood at which minimizes the expectation of the integral over the input space of the variance of an unnormalized posterior density approximation based on a Gaussian process emulator for the log-likelihood function, with the expectation being over the posterior predictive distribution on the unnormalized target density under the Gaussian process.

Parameters:

Name	Type	Description	Default
gp_mean_and_variance	PosteriorPredictiveMeanAndVariance	Function evaluating mean and variance of Gaussian process emulator for target log-density on input space.	required
gp_lookahead_variance_reduction	PosteriorPredictiveLookaheadVarianceReduction	Function evaluating reduction in variance of Gaussian process emulator for log-density at a test-point given one or 'pending' input points, when outputs associated with pending inputs are assumed to follow the posterior predictive distribution under the Gaussian process.	required
integration_inputs	ArrayLike	Points to use when approximating integrals over input space.	required
integration_log_weights	ArrayLike	Logarithm of weights associated with each of points in integration_inputs.	required

Returns:

Type	Description
AcquisitionFunction	The acquisition function to minimize to select new input point(s).

Source code in src/calibr/acquisition_functions.py

def get_expected_integrated_variance_acquisition_function(
    gp_mean_and_variance: PosteriorPredictiveMeanAndVariance,
    gp_lookahead_variance_reduction: PosteriorPredictiveLookaheadVarianceReduction,
    integration_inputs: ArrayLike,
    integration_log_weights: ArrayLike,
) -> AcquisitionFunction:
    """
    Construct acquisition function for minimising expected integrated variance.

    Selects next input points to evaluate log-likelihood at which minimizes the
    expectation of the integral over the input space of the variance of an unnormalized
    posterior density approximation based on a Gaussian process emulator for the
    log-likelihood function, with the expectation being over the posterior predictive
    distribution on the unnormalized target density under the Gaussian process.

    Args:
        gp_mean_and_variance: Function evaluating mean and variance of Gaussian process
            emulator for target log-density on input space.
        gp_lookahead_variance_reduction: Function evaluating reduction in variance of
            Gaussian process emulator for log-density at a test-point given one or
            'pending' input points, when outputs associated with pending inputs are
            assumed to follow the posterior predictive distribution under the Gaussian
            process.
        integration_inputs: Points to use when approximating integrals over input space.
        integration_log_weights: Logarithm of weights associated with each of points
            in `integration_inputs`.

    Returns:
        The acquisition function to minimize to select new input point(s).
    """
    mean_integration_inputs, variance_integration_inputs = jax.vmap(
        gp_mean_and_variance
    )(integration_inputs)

    def acquisition_function(new_inputs: ArrayLike) -> Array:
        lookahead_variance_reduction_integration_inputs = jax.vmap(
            gp_lookahead_variance_reduction, (0, None)
        )(integration_inputs, new_inputs)
        # We neglect the initial constant wrt θ* term in
        # Lᵛₜ(θ*) = ∫ exp(2mₜ(θ) + s²ₜ(θ)) (exp(s²ₜ(θ)) - exp(τ²ₜ(θ; θ*))) dθ
        # and use
        # -log ∫ exp(2mₜ(θ) + s²ₜ(θ) + τ²ₜ(θ; θ*)) dθ
        # corresponding to the negative logarithm of the negation of the second term
        # in the expected integrated variance design criterion.
        # This appears to give a more numerically stable objective function.
        return -jsp.special.logsumexp(
            integration_log_weights
            + 2 * mean_integration_inputs
            + variance_integration_inputs
            + lookahead_variance_reduction_integration_inputs
        )

    return acquisition_function

get_integrated_median_interquantile_range_acquisition_function

get_integrated_median_interquantile_range_acquisition_function(
    gp_mean_and_variance,
    gp_lookahead_variance_reduction,
    integration_inputs,
    integration_log_weights,
    quantile_interval=(0.25, 0.75),
)

Construct acquisition function for minimising integrated median interquantile range.

Selects next input points to evaluate target log-density at which minimizes the integral over the input space of the median interquantile range of an unnormalized target density approximation based on a Gaussian process emulator for the log-density function, with the median being over the posterior predictive distribution on the unnormalized target density under the Gaussian process.

Parameters:

Name	Type	Description	Default
gp_mean_and_variance	PosteriorPredictiveMeanAndVariance	Function evaluating mean and variance of Gaussian process emulator for target log-density on input space.	required
gp_lookahead_variance_reduction	PosteriorPredictiveLookaheadVarianceReduction	Function evaluating reduction in variance of Gaussian process emulator for log-density at a test-point given one or 'pending' input points, when outputs associated with pending inputs are assumed to follow the posterior predictive distribution under the Gaussian process.	required
integration_inputs	ArrayLike	Points to use when approximate integrals over input space.	required
integration_log_weights	ArrayLike	Logarithm of weights associated with each of points in integration_inputs.	required
quantile_interval	tuple[float, float]	Lower and upper quantiles specifying inter-quantile range to optimize.	(0.25, 0.75)

Returns:

Type	Description
AcquisitionFunction	The acquisition function to minimize to select new input point(s).

Source code in src/calibr/acquisition_functions.py

def get_integrated_median_interquantile_range_acquisition_function(
    gp_mean_and_variance: PosteriorPredictiveMeanAndVariance,
    gp_lookahead_variance_reduction: PosteriorPredictiveLookaheadVarianceReduction,
    integration_inputs: ArrayLike,
    integration_log_weights: ArrayLike,
    quantile_interval: tuple[float, float] = (0.25, 0.75),
) -> AcquisitionFunction:
    """
    Construct acquisition function for minimising integrated median interquantile range.

    Selects next input points to evaluate target log-density at which minimizes the
    integral over the input space of the median interquantile range of an unnormalized
    target density approximation based on a Gaussian process emulator for the
    log-density function, with the median being over the posterior predictive
    distribution on the unnormalized target density under the Gaussian process.

    Args:
        gp_mean_and_variance: Function evaluating mean and variance of Gaussian process
            emulator for target log-density on input space.
        gp_lookahead_variance_reduction: Function evaluating reduction in variance of
            Gaussian process emulator for log-density at a test-point given one or
            'pending' input points, when outputs associated with pending inputs are
            assumed to follow the posterior predictive distribution under the Gaussian
            process.
        integration_inputs: Points to use when approximate integrals over input space.
        integration_log_weights: Logarithm of weights associated with each of points
            in `integration_inputs`.
        quantile_interval: Lower and upper quantiles specifying inter-quantile range
            to optimize.

    Returns:
        The acquisition function to minimize to select new input point(s).
    """
    lower = jsp.special.ndtri(quantile_interval[0])
    upper = jsp.special.ndtri(quantile_interval[1])
    mean_integration_inputs, variance_integration_inputs = jax.vmap(
        gp_mean_and_variance
    )(integration_inputs)

    def acquisition_function(new_inputs: ArrayLike) -> Array:
        lookahead_variance_reduction_integration_inputs = jax.vmap(
            gp_lookahead_variance_reduction, (0, None)
        )(integration_inputs, new_inputs)
        lookahead_standard_deviation_integration_inputs = (
            abs(
                variance_integration_inputs
                - lookahead_variance_reduction_integration_inputs
            )
            ** 0.5
        )
        return jsp.special.logsumexp(
            integration_log_weights
            + mean_integration_inputs
            + upper * lookahead_standard_deviation_integration_inputs
            + jnp.log1p(
                -jnp.exp(
                    (lower - upper) * lookahead_standard_deviation_integration_inputs
                )
            )
        )

    return acquisition_function

get_maximum_interquantile_range_greedy_batch_acquisition_functions

get_maximum_interquantile_range_greedy_batch_acquisition_functions(
    gp_mean_and_variance,
    gp_lookahead_variance_reduction,
    quantile_interval=(0.25, 0.75),
)

Construct acquisition function for greedy maximisation of interquantile range.

Selects next input points to evaluate target log-density at which maximise interquantile range of an unnormalized target density approximation based on a Gaussian process emulator for the log-density function. After an initial point is chosen, subsequent points in the batch are selected in a greedy fashion by maximising the interquantile range given the already selected point(s), assuming the log-density values at the already selected points are distributed according the Gaussian process predictive distribution.

Parameters:

Name	Type	Description	Default
gp_mean_and_variance	PosteriorPredictiveMeanAndVariance	Function evaluating mean and variance of Gaussian process emulator for target log-density on input space.	required
gp_lookahead_variance_reduction	PosteriorPredictiveLookaheadVarianceReduction	Function evaluating reduction in variance of Gaussian process emulator for log-density at a test-point given one or 'pending' input points, when outputs associated with pending inputs are assumed to follow the posterior predictive distribution under the Gaussian process.	required
quantile_interval	tuple[float, float]	Lower and upper quantiles specifying inter-quantile range to optimize.	(0.25, 0.75)

Returns:

Type	Description
Callable	The acquisition function to minimize to select new input point(s). The function
Callable	takes one or two arguments. If a single argument is passed it corresponds to
Callable	the acquisition function for the initial point in a batch. If two arguments are
Callable	passed, the first argument corresponds to the new input point being chosen and
Callable	the second argument to the already selected input point(s) in the batch.

Source code in src/calibr/acquisition_functions.py

def get_maximum_interquantile_range_greedy_batch_acquisition_functions(
    gp_mean_and_variance: PosteriorPredictiveMeanAndVariance,
    gp_lookahead_variance_reduction: PosteriorPredictiveLookaheadVarianceReduction,
    quantile_interval: tuple[float, float] = (0.25, 0.75),
) -> Callable:
    """
    Construct acquisition function for greedy maximisation of interquantile range.

    Selects next input points to evaluate target log-density at which maximise
    interquantile range of an unnormalized target density approximation based on a
    Gaussian process emulator for the log-density function. After an initial point is
    chosen, subsequent points in the batch are selected in a greedy fashion by
    maximising the interquantile range given the already selected point(s), assuming the
    log-density values at the already selected points are distributed according the
    Gaussian process predictive distribution.

    Args:
        gp_mean_and_variance: Function evaluating mean and variance of Gaussian process
            emulator for target log-density on input space.
        gp_lookahead_variance_reduction: Function evaluating reduction in variance of
            Gaussian process emulator for log-density at a test-point given one or
            'pending' input points, when outputs associated with pending inputs are
            assumed to follow the posterior predictive distribution under the Gaussian
            process.
        quantile_interval: Lower and upper quantiles specifying inter-quantile range
            to optimize.

    Returns:
        The acquisition function to minimize to select new input point(s). The function
        takes one or two arguments. If a single argument is passed it corresponds to
        the acquisition function for the initial point in a batch. If two arguments are
        passed, the first argument corresponds to the new input point being chosen and
        the second argument to the already selected input point(s) in the batch.
    """
    lower = jsp.special.ndtri(quantile_interval[0])
    upper = jsp.special.ndtri(quantile_interval[1])

    def acquisition_function(
        new_input: ArrayLike, pending_inputs: ArrayLike | None = None
    ) -> Array:
        mean, variance = gp_mean_and_variance(new_input)
        if pending_inputs is None:
            lookahead_variance_reduction = 0
        else:
            lookahead_variance_reduction = gp_lookahead_variance_reduction(
                new_input, pending_inputs
            )
        lookahead_standard_deviation = (
            abs(variance - lookahead_variance_reduction) ** 0.5
        )
        return (
            -mean
            - upper * lookahead_standard_deviation
            - jnp.log1p(-jnp.exp(lookahead_standard_deviation * (lower - upper)))
        )

    return acquisition_function

get_maximum_variance_greedy_batch_acquisition_functions

get_maximum_variance_greedy_batch_acquisition_functions(
    gp_mean_and_variance, gp_lookahead_variance_reduction
)

Construct acquisition functions for greedy maximisation of variance.

Selects next input points to evaluate target log-density at which maximise variance of an unnormalized target density approximation based on a Gaussian process emulator for the log-density function. After an initial point is chosen, subsequent points in the batch are selected in a greedy fashion by maximising the variance given the already selected point(s), assuming the log-density values at the already selected points are distributed according the Gaussian process predictive distribution.

Parameters:

Name	Type	Description	Default
gp_mean_and_variance	PosteriorPredictiveMeanAndVariance	Function evaluating mean and variance of Gaussian process emulator for target log-density on input space.	required
gp_lookahead_variance_reduction	PosteriorPredictiveLookaheadVarianceReduction	Function evaluating reduction in variance of Gaussian process emulator for log-density at a test-point given one or 'pending' input points, when outputs associated with pending inputs are assumed to follow the posterior predictive distribution under the Gaussian process.	required

Returns:

Type	Description
Callable	The acquisition function to minimize to select new input point(s). The function
Callable	takes one or two arguments. If a single argument is passed it corresponds to
Callable	the acquisition function for the initial point in a batch. If two arguments are
Callable	passed, the first argument corresponds to the new input point being chosen and
Callable	the second argument to the already selected input point(s) in the batch.

Source code in src/calibr/acquisition_functions.py

def get_maximum_variance_greedy_batch_acquisition_functions(
    gp_mean_and_variance: PosteriorPredictiveMeanAndVariance,
    gp_lookahead_variance_reduction: PosteriorPredictiveLookaheadVarianceReduction,
) -> Callable:
    """
    Construct acquisition functions for greedy maximisation of variance.

    Selects next input points to evaluate target log-density at which maximise variance
    of an unnormalized target density approximation based on a Gaussian process emulator
    for the log-density function. After an initial point is chosen, subsequent points in
    the batch are selected in a greedy fashion by maximising the variance given the
    already selected point(s), assuming the log-density values at the already selected
    points are distributed according the Gaussian process predictive distribution.

    Args:
        gp_mean_and_variance: Function evaluating mean and variance of Gaussian process
            emulator for target log-density on input space.
        gp_lookahead_variance_reduction: Function evaluating reduction in variance of
            Gaussian process emulator for log-density at a test-point given one or
            'pending' input points, when outputs associated with pending inputs are
            assumed to follow the posterior predictive distribution under the Gaussian
            process.

    Returns:
        The acquisition function to minimize to select new input point(s). The function
        takes one or two arguments. If a single argument is passed it corresponds to
        the acquisition function for the initial point in a batch. If two arguments are
        passed, the first argument corresponds to the new input point being chosen and
        the second argument to the already selected input point(s) in the batch.
    """

    def acquisition_function(
        new_input: ArrayLike, pending_inputs: ArrayLike | None = None
    ) -> Array:
        mean, variance = gp_mean_and_variance(new_input)
        if pending_inputs is None:
            lookahead_variance_reduction = 0
        else:
            lookahead_variance_reduction = gp_lookahead_variance_reduction(
                new_input, pending_inputs
            )
        lookahead_variance = variance - lookahead_variance_reduction
        return -2 * (mean + variance) - jnp.log1p(-jnp.exp(-lookahead_variance))

    return acquisition_function

calibration

Functions for iteratively calibrating the parameters of a probabilistic model.

calibrate

calibrate(
    num_initial_inputs,
    batch_size,
    num_iterations,
    rng,
    sample_initial_inputs,
    posterior_log_density_batch,
    gaussian_process_factory,
    get_integration_points_and_log_weights,
    *,
    fit_gaussian_process_parameters=fit_gaussian_process_parameters_map,
    get_acquisition_function=get_integrated_median_interquantile_range_acquisition_function,
    get_next_inputs_batch=get_next_inputs_batch_by_joint_optimization,
    end_of_iteration_callback=None
)

Estimate the posterior on the unknown inputs of an (expensive to evaluate) model.

Iterates evaluating log density for posterior at a batch of inputs (unknown variables to infer posterior on), fitting a Gaussian process to the log density evaluations so far and optimizing an acquisition function using the Gaussian process emulator to choose a new batch of input points at which to evaluate the log density (by minimizing a measure of the expected uncertainty in the emulator about the log posterior density function).

Parameters:

Name	Type	Description	Default
num_initial_inputs	int	Number of initial inputs to evaluate posterior log density at to initialize calibration.	required
batch_size	int	Size of batch of inputs to optimize for and evaluate log density at in each calibration iteration.	required
num_iterations	int	Number of calibration iterations to perform. Total number of model posterior log density evaluations is num_initial_inputs + batch_size * num_iterations.	required
rng	Generator	Seeded NumPy random number generator.	required
sample_initial_inputs	InitialInputSampler	Function outputting reasonable random initial values for batch of inputs when passed a random number generator and batch size.	required
posterior_log_density_batch	Callable[[ArrayLike], Array]	Function computing logarithm of (unnormalized) posterior density on model inputs, for a batch of inputs (with passed argument being a two dimensional array with first dimension the batch index).	required
gaussian_process_factory	GaussianProcessFactory	Factory function generating Gaussian process models given a data dictionary.	required
get_integration_points_and_log_weights	Callable[[Generator, InitialInputSampler, PosteriorPredictiveMeanAndVariance], tuple[Array, Array]]	Function which outputs points in input space and corresponding (log) weights by which to estimate integrals over the input space in the acquisition function as a weighted sum. The function is passed a seeded random number generator, a function to sample random points in the input space and the current Gaussian process posterior predictive mean and variance function. The input points and weights may for example be generated according to a (deterministic) numerical quadrature rule for low dimensionalities, a stochastic (quasi-) Monte Carlo scheme for moderate dimensionalities or a Markov chain Monte Carlo or sequential Monte Carlo scheme for higher dimensionalities.	required
fit_gaussian_process_parameters	GaussianProcessParameterFitter	Function which fits the parameters of a Gaussian process model given the current data (input-output pairs). Passed a seeded random number generator and tuple of Gaussian process model functions.	fit_gaussian_process_parameters_map
get_acquisition_function	AcquisitionFunctionFactory	Factory function generating acquisition functions given Gaussian process posterior predictive functions.	get_integrated_median_interquantile_range_acquisition_function
get_next_inputs_batch	Callable[[Generator, AcquisitionFunction, InitialInputSampler, int], tuple[Array, float]]	Function which computes next batch of inputs to evaluate model at by optimizing the current acquisition function. Passed a seeded random number generator, acquisition function, input sampler and batch size.	get_next_inputs_batch_by_joint_optimization
end_of_iteration_callback	EndOfIterationCallback | None	Optional callback function evaluate at end of each calibration iteration, for example for logging metrics or plotting / saving intermediate outputs. Passed current iteration index, Gaussian process posterior mean and variance, data dictionary with all inputs and corresponding log density evaluations so far, batch of inputs selected in current iteration and corresponding optimized acquisition function value.	None

Returns:

Type	Description
GaussianProcessModel	Tuple of Gaussian process model, data dictionary containing all model inputs
DataDict	and log density evaluations and fitted Gaussian process model parameters at
ParametersDict	final iteration.

Source code in src/calibr/calibration.py

def calibrate(  # noqa: PLR0913
    num_initial_inputs: int,
    batch_size: int,
    num_iterations: int,
    rng: Generator,
    sample_initial_inputs: InitialInputSampler,
    posterior_log_density_batch: Callable[[ArrayLike], Array],
    gaussian_process_factory: GaussianProcessFactory,
    get_integration_points_and_log_weights: Callable[
        [Generator, InitialInputSampler, PosteriorPredictiveMeanAndVariance],
        tuple[Array, Array],
    ],
    *,
    fit_gaussian_process_parameters: GaussianProcessParameterFitter = (
        fit_gaussian_process_parameters_map
    ),
    get_acquisition_function: AcquisitionFunctionFactory = (
        get_integrated_median_interquantile_range_acquisition_function
    ),
    get_next_inputs_batch: Callable[
        [Generator, AcquisitionFunction, InitialInputSampler, int], tuple[Array, float]
    ] = get_next_inputs_batch_by_joint_optimization,
    end_of_iteration_callback: EndOfIterationCallback | None = None,
) -> tuple[GaussianProcessModel, DataDict, ParametersDict]:
    """
    Estimate the posterior on the unknown inputs of an (expensive to evaluate) model.

    Iterates evaluating log density for posterior at a batch of inputs (unknown
    variables to infer posterior on), fitting a Gaussian process to the log density
    evaluations so far and optimizing an acquisition function using the Gaussian process
    emulator to choose a new batch of input points at which to evaluate the log density
    (by minimizing a measure of the expected uncertainty in the emulator about the log
    posterior density function).

    Args:
        num_initial_inputs: Number of initial inputs to evaluate posterior log density
            at to initialize calibration.
        batch_size: Size of batch of inputs to optimize for and evaluate log density at
            in each calibration iteration.
        num_iterations: Number of calibration iterations to perform. Total number of
            model posterior log density evaluations is
            `num_initial_inputs + batch_size * num_iterations`.
        rng: Seeded NumPy random number generator.
        sample_initial_inputs: Function outputting reasonable random initial values for
            batch of inputs when passed a random number generator and batch size.
        posterior_log_density_batch: Function computing logarithm of (unnormalized)
            posterior density on model inputs, for a batch of inputs (with passed
            argument being a two dimensional array with first dimension the batch
            index).
        gaussian_process_factory: Factory function generating Gaussian process models
            given a data dictionary.
        get_integration_points_and_log_weights: Function which outputs points in input
            space and corresponding (log) weights by which to estimate integrals over
            the input space in the acquisition function as a weighted sum. The function
            is passed a seeded random number generator, a function to sample random
            points in the input space and the current Gaussian process posterior
            predictive mean and variance function. The input points and weights may for
            example be generated according to a (deterministic) numerical quadrature
            rule for low dimensionalities, a stochastic (quasi-) Monte Carlo scheme for
            moderate dimensionalities or a Markov chain Monte Carlo or sequential Monte
            Carlo scheme for higher dimensionalities.
        fit_gaussian_process_parameters: Function which fits the parameters of a
            Gaussian process model given the current data (input-output pairs). Passed
            a seeded random number generator and tuple of Gaussian process model
            functions.
        get_acquisition_function: Factory function generating acquisition functions
            given Gaussian process posterior predictive functions.
        get_next_inputs_batch: Function which computes next batch of inputs to evaluate
            model at by optimizing the current acquisition function. Passed a seeded
            random number generator, acquisition function, input sampler and batch size.
        end_of_iteration_callback: Optional callback function evaluate at end of each
            calibration iteration, for example for logging metrics or plotting / saving
            intermediate outputs. Passed current iteration index, Gaussian process
            posterior mean and variance, data dictionary with all inputs and
            corresponding log density evaluations so far, batch of inputs selected in
            current iteration and corresponding optimized acquisition function value.

    Returns:
        Tuple of Gaussian process model, data dictionary containing all model inputs
        and log density evaluations and fitted Gaussian process model parameters at
        final iteration.
    """
    inputs = sample_initial_inputs(rng, num_initial_inputs)
    data = {"inputs": inputs, "outputs": posterior_log_density_batch(inputs)}
    for iteration_index in range(num_iterations + 1):
        gaussian_process = gaussian_process_factory(data)
        parameters = fit_gaussian_process_parameters(rng, gaussian_process)
        (
            posterior_mean_and_variance,
            lookahead_variance_reduction,
        ) = gaussian_process.get_posterior_functions(parameters)
        if iteration_index == num_iterations:
            # In final iteration, only fit Gaussian process to all model evaluations
            # computed so far without acquiring a new set of inputs to evaluate
            break
        (
            integration_inputs,
            integration_log_weights,
        ) = get_integration_points_and_log_weights(
            rng, sample_initial_inputs, posterior_mean_and_variance
        )
        acquisition_function = get_acquisition_function(
            posterior_mean_and_variance,
            lookahead_variance_reduction,
            integration_inputs,
            integration_log_weights,
        )
        next_inputs, acquisition_function_value = get_next_inputs_batch(
            rng, acquisition_function, sample_initial_inputs, batch_size
        )
        next_outputs = posterior_log_density_batch(next_inputs)
        data = {
            "inputs": np.concatenate((data["inputs"], next_inputs)),
            "outputs": np.concatenate((data["outputs"], next_outputs)),
        }
        if end_of_iteration_callback is not None:
            end_of_iteration_callback(
                iteration_index,
                posterior_mean_and_variance,
                data,
                next_inputs,
                acquisition_function_value,
            )
    return gaussian_process, data, parameters

get_next_inputs_batch_by_greedy_optimization

get_next_inputs_batch_by_greedy_optimization(
    rng,
    acquisition_function,
    sample_initial_inputs,
    batch_size,
    *,
    minimize_function=minimize_with_restarts,
    **minimize_function_kwargs
)

Get next batch of inputs to evaluate by greedily optimizing acquisition function.

Sequentially minimizes acquisition function for b in 1 to batch_size by fixing b - 1 inputs already optimized and minimizing over a single new input in each iteration.

Parameters:

Name	Type	Description	Default
rng	Generator	NumPy random number generator for initializing optimization runs.	required
acquisition_function	AcquisitionFunction	Scalar-valued function of a batch of inputs to optimize to find new batch of inputs to evaluate model for.	required
sample_initial_inputs	InitialInputSampler	Function outputting reasonable random initial values for batch of inputs when passed a random number generator and batch size. Used to initialize state for optimization runs.	required
batch_size	int	Number of inputs in batch.	required
minimize_function	GlobalMinimizer	Function used to attempt to find minimum of (sequence of) acquisition functions.	minimize_with_restarts
**minimize_function_kwargs	GlobalMinimizerKwarg	Any keyword arguments to pass to minimize_function function used to optimize acquisition function.	{}

Returns:

Type	Description
tuple[Array, float]	Tuple of optimized inputs batch and corresponding value of acquisition function.

Source code in src/calibr/calibration.py

def get_next_inputs_batch_by_greedy_optimization(
    rng: Generator,
    acquisition_function: AcquisitionFunction,
    sample_initial_inputs: InitialInputSampler,
    batch_size: int,
    *,
    minimize_function: GlobalMinimizer = minimize_with_restarts,
    **minimize_function_kwargs: GlobalMinimizerKwarg,
) -> tuple[Array, float]:
    """
    Get next batch of inputs to evaluate by greedily optimizing acquisition function.

    Sequentially minimizes acquisition function for `b` in 1 to `batch_size` by fixing
    `b - 1` inputs already optimized and minimizing over a single new input in each
    iteration.

    Args:
        rng: NumPy random number generator for initializing optimization runs.
        acquisition_function: Scalar-valued function of a batch of inputs to optimize to
            find new batch of inputs to evaluate model for.
        sample_initial_inputs: Function outputting reasonable random initial values for
            batch of inputs when passed a random number generator and batch size. Used
            to initialize state for optimization runs.
        batch_size: Number of inputs in batch.
        minimize_function: Function used to attempt to find minimum of (sequence of)
            acquisition functions.
        **minimize_function_kwargs: Any keyword arguments to pass to
            `minimize_function` function used to optimize acquisition function.

    Returns:
        Tuple of optimized inputs batch and corresponding value of acquisition function.
    """

    def acquisition_function_greedy(
        current_input: ArrayLike, fixed_inputs: list[ArrayLike]
    ) -> float:
        return acquisition_function(jnp.stack([current_input, *fixed_inputs]))

    fixed_inputs: list[ArrayLike] = []
    for _ in range(batch_size):
        current_input, min_acquisition_function = minimize_function(
            objective_function=partial(
                acquisition_function_greedy, fixed_inputs=fixed_inputs
            ),
            sample_initial_state=lambda r: sample_initial_inputs(r, 1).flatten(),
            rng=rng,
            **minimize_function_kwargs,
        )
        fixed_inputs.append(current_input)

    return np.stack(fixed_inputs), min_acquisition_function

get_next_inputs_batch_by_joint_optimization

get_next_inputs_batch_by_joint_optimization(
    rng,
    acquisition_function,
    sample_initial_inputs,
    batch_size,
    *,
    minimize_function=minimize_with_restarts,
    **minimize_function_kwargs
)

Get next batch of inputs to evaluate by jointly optimizing acquisition function.

Minimizes acquisition function over product of batch_size input spaces.

Parameters:

Name	Type	Description	Default
rng	Generator	NumPy random number generator for initializing optimization runs.	required
acquisition_function	AcquisitionFunction	Scalar-valued function of a batch of inputs to optimize to find new batch of inputs to evaluate model for.	required
sample_initial_inputs	InitialInputSampler	Function outputting reasonable random initial values for batch of inputs when passed a random number generator and batch size. Used to initialize state for optimization runs.	required
batch_size	int	Number of inputs in batch.	required
minimize_function	GlobalMinimizer	Function used to attempt to find minimum of acquisition function.	minimize_with_restarts
**minimize_function_kwargs	GlobalMinimizerKwarg	Any keyword arguments to pass to minimize_function function used to optimize acquisition function.	{}

Returns:

Type	Description
tuple[Array, float]	Tuple of optimized inputs batch and corresponding value of acquisition function.

Source code in src/calibr/calibration.py

def get_next_inputs_batch_by_joint_optimization(
    rng: Generator,
    acquisition_function: AcquisitionFunction,
    sample_initial_inputs: InitialInputSampler,
    batch_size: int,
    *,
    minimize_function: GlobalMinimizer = minimize_with_restarts,
    **minimize_function_kwargs: GlobalMinimizerKwarg,
) -> tuple[Array, float]:
    """
    Get next batch of inputs to evaluate by jointly optimizing acquisition function.

    Minimizes acquisition function over product of `batch_size` input spaces.

    Args:
        rng: NumPy random number generator for initializing optimization runs.
        acquisition_function: Scalar-valued function of a batch of inputs to optimize to
            find new batch of inputs to evaluate model for.
        sample_initial_inputs: Function outputting reasonable random initial values for
            batch of inputs when passed a random number generator and batch size. Used
            to initialize state for optimization runs.
        batch_size: Number of inputs in batch.
        minimize_function: Function used to attempt to find minimum of acquisition
            function.
        **minimize_function_kwargs: Any keyword arguments to pass to
            `minimize_function` function used to optimize acquisition function.

    Returns:
        Tuple of optimized inputs batch and corresponding value of acquisition function.
    """

    def acquisition_function_flat_input(flat_inputs: ArrayLike) -> float:
        return acquisition_function(flat_inputs.reshape((batch_size, -1)))

    if minimize_function is minimize_with_restarts:
        minimize_function_kwargs.setdefault("number_minima_to_find", 5)
        minimize_function_kwargs.setdefault("maximum_minimize_calls", 100)
        minimize_function_kwargs.setdefault("minimize_method", "Newton-CG")

    flat_inputs, min_acquisition_function = minimize_function(
        objective_function=acquisition_function_flat_input,
        sample_initial_state=lambda r: sample_initial_inputs(r, batch_size).flatten(),
        rng=rng,
        **minimize_function_kwargs,
    )

    return flat_inputs.reshape((batch_size, -1)), min_acquisition_function

emulation

Functions and types for constructing and fitting Gaussian process emulators.

GaussianProcessModel

Bases: NamedTuple

Wrapper for functions associated with a Gaussian process model.

Source code in src/calibr/emulation.py

class GaussianProcessModel(NamedTuple):
    """Wrapper for functions associated with a Gaussian process model."""

    neg_log_marginal_posterior: Callable[[ArrayLike], float]
    get_posterior_functions: PosteriorPredictiveFunctionFactory
    transform_parameters: ParameterTransformer
    sample_unconstrained_parameters: UnconstrainedParametersSampler

fit_gaussian_process_parameters_hmc

fit_gaussian_process_parameters_hmc(
    rng,
    gaussian_process,
    *,
    n_chain=1,
    n_warm_up_iter=500,
    n_main_iter=1000,
    r_hat_threshold=None
)

Fit parameters of Gaussian process model by sampling posterior using HMC.

Uses Hamiltonian Monte Carlo (HMC) to generate chain(s) of samples approximating posterior distribution on Gaussian process parameters given data.

Parameters:

Name	Type	Description	Default
rng	Generator	Seeded NumPy random number generator.	required
gaussian_process	GaussianProcessModel	Tuple of functions for Gaussian process model to fit.	required
n_chain	int	Number of Markov chains to simulate.	1
n_warm_up_iter	int	Number of adaptive warm-up iterations to run for each chain.	500
n_main_iter	int	Number of main sampling stage iterations to run for each chain.	1000
r_hat_threshold	float | None	If not None, specifies a maximum value for the (rank-normalized, split) R-hat convergence diagnostic computed from the chains, with R-hat values exceeding this threshold leading to an exception being raised. Requires n_chain > 1 and for ArviZ package to be installed.	None

Returns:

Type	Description
ParametersDict	Dictionary of parameters corresponding to approximate posterior sample.

Source code in src/calibr/emulation.py

def fit_gaussian_process_parameters_hmc(
    rng: Generator,
    gaussian_process: GaussianProcessModel,
    *,
    n_chain: int = 1,
    n_warm_up_iter: int = 500,
    n_main_iter: int = 1000,
    r_hat_threshold: float | None = None,
) -> ParametersDict:
    """
    Fit parameters of Gaussian process model by sampling posterior using HMC.

    Uses Hamiltonian Monte Carlo (HMC) to generate chain(s) of samples approximating
    posterior distribution on Gaussian process parameters given data.

    Args:
        rng: Seeded NumPy random number generator.
        gaussian_process: Tuple of functions for Gaussian process model to fit.
        n_chain: Number of Markov chains to simulate.
        n_warm_up_iter: Number of adaptive warm-up iterations to run for each chain.
        n_main_iter: Number of main sampling stage iterations to run for each chain.
        r_hat_threshold: If not `None`, specifies a maximum value for the
            (rank-normalized, split) R-hat convergence diagnostic computed from the
            chains, with R-hat values exceeding this threshold leading to an
            exception being raised. Requires `n_chain > 1` and for ArviZ package to
            be installed.

    Returns:
        Dictionary of parameters corresponding to approximate posterior sample.
    """
    if r_hat_threshold is not None and not ARVIZ_IMPORTED:
        msg = "R-hat convergence checks require ArviZ to be installed"
        raise RuntimeError(msg)
    value_and_grad_neg_log_marginal_posterior = jax.jit(
        jax.value_and_grad(gaussian_process.neg_log_marginal_posterior)
    )

    def grad_neg_log_marginal_posterior(
        unconstrained_variables: ArrayLike,
    ) -> tuple[Array, float]:
        value, grad = value_and_grad_neg_log_marginal_posterior(
            unconstrained_variables
        )
        return np.asarray(grad), float(value)

    init_states = gaussian_process.sample_unconstrained_parameters(rng, n_chain)

    system = mici.systems.EuclideanMetricSystem(
        neg_log_dens=gaussian_process.neg_log_marginal_posterior,
        grad_neg_log_dens=grad_neg_log_marginal_posterior,
    )
    integrator = mici.integrators.LeapfrogIntegrator(system)
    sampler = mici.samplers.DynamicMultinomialHMC(system, integrator, rng)

    final_states, traces, _ = sampler.sample_chains(
        n_warm_up_iter,
        n_main_iter,
        init_states,
        monitor_stats=["accept_stat", "step_size", "n_step", "diverging"],
        adapters=[
            mici.adapters.DualAveragingStepSizeAdapter(0.8),
            mici.adapters.OnlineCovarianceMetricAdapter(),
        ],
        n_process=1,
    )

    if n_chain > 1 and r_hat_threshold is not None:
        max_rhat = float(arviz.rhat(traces).to_array().max())
        if max_rhat > r_hat_threshold:
            msg = f"Chain convergence issue: max rank-normalized R-hat {max_rhat}"
            raise RuntimeError(msg)
    return gaussian_process.transform_parameters(final_states[0].pos)

fit_gaussian_process_parameters_map

fit_gaussian_process_parameters_map(
    rng,
    gaussian_process,
    *,
    minimize_function=minimize_with_restarts,
    **minimize_function_kwargs
)

Fit parameters of Gaussian process model by maximimizing posterior density.

Finds maximum-a-posterior (MAP) estimate of Gaussian process parameters by minimizing negative logarithm of posterior density on parameters given data.

Parameters:

Name	Type	Description	Default
rng	Generator	Seeded NumPy random number generator.	required
gaussian_process	GaussianProcessModel	Tuple of functions for Gaussian process model to fit.	required
minimize_function	GlobalMinimizer	Function used to attempt to find global minimum of negative log posterior density function.	minimize_with_restarts
**minimize_function_kwargs	GlobalMinimizerKwarg	Any keyword arguments to pass to minimize_function function used to optimize negative posterior log density function.	{}

Returns:

Type	Description
ParametersDict	Dictionary of parameters corresponding to maximum-a-posteriori estimate.

Source code in src/calibr/emulation.py

def fit_gaussian_process_parameters_map(
    rng: Generator,
    gaussian_process: GaussianProcessModel,
    *,
    minimize_function: GlobalMinimizer = minimize_with_restarts,
    **minimize_function_kwargs: GlobalMinimizerKwarg,
) -> ParametersDict:
    """Fit parameters of Gaussian process model by maximimizing posterior density.

    Finds maximum-a-posterior (MAP) estimate of Gaussian process parameters by
    minimizing negative logarithm of posterior density on parameters given data.

    Args:
        rng: Seeded NumPy random number generator.
        gaussian_process: Tuple of functions for Gaussian process model to fit.
        minimize_function: Function used to attempt to find global minimum of negative
            log posterior density function.
        **minimize_function_kwargs: Any keyword arguments to pass to
            `minimize_function` function used to optimize negative posterior log density
            function.

    Returns:
        Dictionary of parameters corresponding to maximum-a-posteriori estimate.
    """
    if minimize_function is minimize_with_restarts:
        minimize_function_kwargs.setdefault("number_minima_to_find", 1)
        minimize_function_kwargs.setdefault("maximum_minimize_calls", 100)
        minimize_function_kwargs.setdefault("minimize_method", "Newton-CG")
    unconstrained_parameters, _ = minimize_function(
        objective_function=gaussian_process.neg_log_marginal_posterior,
        sample_initial_state=lambda r: gaussian_process.sample_unconstrained_parameters(
            r,
            None,
        ),
        rng=rng,
        **minimize_function_kwargs,
    )
    return gaussian_process.transform_parameters(unconstrained_parameters)

get_gaussian_process_factory

get_gaussian_process_factory(
    mean_function,
    covariance_function,
    neg_log_prior_density,
    transform_parameters,
    sample_unconstrained_parameters,
)

Construct a factory function generating Gaussian process models given data.

Parameters:

Name	Type	Description	Default
mean_function	MeanFunction	Mean function for Gaussian process.	required
covariance_function	CovarianceFunction	Covariance function for Gaussian process.	required
neg_log_prior_density	NegativeLogPriorDensity	Negative logarithm of density of prior distribution on vector of unconstrained parameters for Gaussian process model.	required
transform_parameters	ParameterTransformer	Function which maps flat unconstrained parameter vector to a dictionary of (potential constrained) parameters, keyed by parameter name.	required
sample_unconstrained_parameters	UnconstrainedParametersSampler	Function generating random values for unconstrained vector of Gaussian process parameters.	required

Returns:

Type	Description
GaussianProcessFactory	Gaussian process factory function.

Source code in src/calibr/emulation.py

def get_gaussian_process_factory(
    mean_function: MeanFunction,
    covariance_function: CovarianceFunction,
    neg_log_prior_density: NegativeLogPriorDensity,
    transform_parameters: ParameterTransformer,
    sample_unconstrained_parameters: UnconstrainedParametersSampler,
) -> GaussianProcessFactory:
    """Construct a factory function generating Gaussian process models given data.

    Args:
        mean_function: Mean function for Gaussian process.
        covariance_function: Covariance function for Gaussian process.
        neg_log_prior_density: Negative logarithm of density of prior distribution on
            vector of unconstrained parameters for Gaussian process model.
        transform_parameters: Function which maps flat unconstrained parameter vector to
            a dictionary of (potential constrained) parameters, keyed by parameter name.
        sample_unconstrained_parameters: Function generating random values for
            unconstrained vector of Gaussian process parameters.

    Returns:
        Gaussian process factory function.
    """

    def gaussian_process_factory(data: DataDict) -> GaussianProcessModel:
        (
            neg_log_marginal_likelihood,
            get_posterior_functions,
        ) = gaussian_process_with_isotropic_gaussian_observations(
            data, mean_function, covariance_function
        )

        @jax.jit
        def neg_log_marginal_posterior(unconstrained_parameters: ArrayLike) -> float:
            parameters = transform_parameters(unconstrained_parameters)
            return neg_log_prior_density(
                unconstrained_parameters
            ) + neg_log_marginal_likelihood(parameters)

        return GaussianProcessModel(
            neg_log_marginal_posterior,
            get_posterior_functions,
            transform_parameters,
            sample_unconstrained_parameters,
        )

    return gaussian_process_factory

optimization

Functions for minimization of objective functions.

ConvergenceError

Bases: Exception

Error raised when optimizer fails to converge within given computation budget.

Source code in src/calibr/optimization.py

class ConvergenceError(Exception):
    """Error raised when optimizer fails to converge within given computation budget."""

GlobalMinimizer

Bases: Protocol

Function which attempts to find global minimum of a scalar objective function.

Source code in src/calibr/optimization.py

class GlobalMinimizer(Protocol):
    """Function which attempts to find global minimum of a scalar objective function."""

    def __call__(
        self,
        objective_function: ObjectiveFunction,
        sample_initial_state: InitialStateSampler,
        rng: Generator,
        **kwargs: GlobalMinimizerKwarg,
    ) -> tuple[jax.Array, float]:
        """
        Minimize a differentiable objective function.

        Args:
            objective_function: Differentiable scalar-valued function of a single flat
                vector argument to be minimized. Assumed to be specified using JAX
                primitives such that its gradient and Hessian can be computed using
                JAX's automatic differentiation support, and to be suitable for
                just-in-time compilation.
            sample_initial_state: Callable with one argument, which when passed a NumPy
                random number generator returns a random initial state for optimization
                of appropriate dimension.
            rng: Seeded NumPy random number generator.
            **kwargs: Any keyword arguments to global minimizer function.

        Returns:
            Tuple with first entry the state corresponding to the minima point and the
            second entry the corresponding objective function value.
        """

__call__

__call__(objective_function, sample_initial_state, rng, **kwargs)

Minimize a differentiable objective function.

Parameters:

Name	Type	Description	Default
objective_function	ObjectiveFunction	Differentiable scalar-valued function of a single flat vector argument to be minimized. Assumed to be specified using JAX primitives such that its gradient and Hessian can be computed using JAX's automatic differentiation support, and to be suitable for just-in-time compilation.	required
sample_initial_state	InitialStateSampler	Callable with one argument, which when passed a NumPy random number generator returns a random initial state for optimization of appropriate dimension.	required
rng	Generator	Seeded NumPy random number generator.	required
**kwargs	GlobalMinimizerKwarg	Any keyword arguments to global minimizer function.	{}

Returns:

Type	Description
Array	Tuple with first entry the state corresponding to the minima point and the
float	second entry the corresponding objective function value.

Source code in src/calibr/optimization.py

def __call__(
    self,
    objective_function: ObjectiveFunction,
    sample_initial_state: InitialStateSampler,
    rng: Generator,
    **kwargs: GlobalMinimizerKwarg,
) -> tuple[jax.Array, float]:
    """
    Minimize a differentiable objective function.

    Args:
        objective_function: Differentiable scalar-valued function of a single flat
            vector argument to be minimized. Assumed to be specified using JAX
            primitives such that its gradient and Hessian can be computed using
            JAX's automatic differentiation support, and to be suitable for
            just-in-time compilation.
        sample_initial_state: Callable with one argument, which when passed a NumPy
            random number generator returns a random initial state for optimization
            of appropriate dimension.
        rng: Seeded NumPy random number generator.
        **kwargs: Any keyword arguments to global minimizer function.

    Returns:
        Tuple with first entry the state corresponding to the minima point and the
        second entry the corresponding objective function value.
    """

basin_hopping

basin_hopping(
    objective_function,
    sample_initial_state,
    rng,
    *,
    num_iterations=5,
    minimize_method="Newton-CG",
    minimize_max_iterations=None,
    minimize_tol=None,
    **unknown_kwargs
)

Minimize a differentiable objective function with SciPy basin-hopping algorithm.

The basin-hopping algorithm nests an inner local minimization using the scipy.optimize.minimize method within an outer global stepping algorithm which perturbs the current state with a random displacement and accepts or rejects this proposal using a Metropolis criterion.

Parameters:

Name	Type	Description	Default
objective_function	ObjectiveFunction	Differentiable scalar-valued function of a single flat vector argument to be minimized. Assumed to be specified using JAX primitives such that its gradient and Hessian can be computed using JAX's automatic differentiation support, and to be suitable for just-in-time compilation.	required
sample_initial_state	InitialStateSampler	Callable with one argument, which when passed a NumPy random number generator returns a random initial state for optimization of appropriate dimension.	required
rng	Generator	Seeded NumPy random number generator.	required
num_iterations	int	Number of basin-hopping iterations, with number of inner scipy.optimize.minimize calls being num_iterations + 1.	5
minimize_method	str	String specifying one of local optimization methods which can be passed to method argument of scipy.optimize.minimize.	'Newton-CG'
minimize_max_iterations	int | None	Maximum number of iterations in inner local minimization.	None
minimize_tol	float | None	Tolerance parameter for inner local minimization.	None

Returns:

Type	Description
Array	Tuple with first entry the state corresponding to the best minima candidate
float	found and the second entry the corresponding objective function value.

Source code in src/calibr/optimization.py

def basin_hopping(
    objective_function: ObjectiveFunction,
    sample_initial_state: InitialStateSampler,
    rng: Generator,
    *,
    num_iterations: int = 5,
    minimize_method: str = "Newton-CG",
    minimize_max_iterations: int | None = None,
    minimize_tol: float | None = None,
    **unknown_kwargs: GlobalMinimizerKwarg,
) -> tuple[jax.Array, float]:
    """Minimize a differentiable objective function with SciPy basin-hopping algorithm.

    The basin-hopping algorithm nests an inner local minimization using the
    `scipy.optimize.minimize` method within an outer global stepping algorithm which
    perturbs the current state with a random displacement and accepts or rejects this
    proposal using a Metropolis criterion.

    Args:
        objective_function: Differentiable scalar-valued function of a single flat
            vector argument to be minimized. Assumed to be specified using JAX
            primitives such that its gradient and Hessian can be computed using JAX's
            automatic differentiation support, and to be suitable for just-in-time
            compilation.
        sample_initial_state: Callable with one argument, which when passed a NumPy
            random number generator returns a random initial state for optimization of
            appropriate dimension.
        rng: Seeded NumPy random number generator.
        num_iterations: Number of basin-hopping iterations, with number of inner
            `scipy.optimize.minimize` calls being `num_iterations + 1`.
        minimize_method: String specifying one of local optimization methods which can
            be passed to `method` argument of `scipy.optimize.minimize`.
        minimize_max_iterations: Maximum number of iterations in inner local
            minimization.
        minimize_tol: Tolerance parameter for inner local minimization.

    Returns:
        Tuple with first entry the state corresponding to the best minima candidate
        found and the second entry the corresponding objective function value.
    """
    _check_unknown_kwargs(unknown_kwargs)
    results = _basin_hopping(
        jax.jit(objective_function),
        x0=sample_initial_state(rng),
        niter=num_iterations,
        minimizer_kwargs={
            "method": minimize_method,
            "jac": jax.jit(jax.grad(objective_function)),
            "hessp": jax.jit(hessian_vector_product(objective_function)),
            "tol": minimize_tol,
            "options": {"maxiter": minimize_max_iterations},
        },
        seed=rng,
    )
    return results.x, float(results.fun)

hessian_vector_product

hessian_vector_product(scalar_function)

Construct function to compute Hessian-vector product for scalar-valued function.

Parameters:

Name	Type	Description	Default
scalar_function	ObjectiveFunction	Scalar-valued objective function.	required

Returns:

Type	Description
Callable[[ArrayLike, ArrayLike], Array]	Hessian-vector product function.

Source code in src/calibr/optimization.py

def hessian_vector_product(
    scalar_function: ObjectiveFunction,
) -> Callable[[ArrayLike, ArrayLike], jax.Array]:
    """
    Construct function to compute Hessian-vector product for scalar-valued function.

    Args:
        scalar_function: Scalar-valued objective function.

    Returns:
        Hessian-vector product function.
    """

    def hvp(x: ArrayLike, v: ArrayLike) -> jax.Array:
        return jax.jvp(jax.grad(scalar_function), (x,), (v,))[1]

    return hvp

minimize_with_restarts

minimize_with_restarts(
    objective_function,
    sample_initial_state,
    rng,
    *,
    number_minima_to_find=5,
    maximum_minimize_calls=100,
    minimize_method="Newton-CG",
    minimize_max_iterations=None,
    minimize_tol=None,
    logging_function=lambda _: None,
    **unknown_kwargs
)

Minimize a differentiable objective function with random restarts.

Iteratively calls scipy.optimize.minimize to attempt to find a minimum of an objective function until a specified number of candidate minima are successfully found, with the initial state for each minimize called being randomly sampled using a user provided function. The candidate minima with the minimum value for the objective function is returned along with the corresponding objective function value.

Parameters:

Name	Type	Description	Default
objective_function	ObjectiveFunction	Differentiable scalar-valued function of a single flat vector argument to be minimized. Assumed to be specified using JAX primitives such that its gradient and Hessian can be computed using JAX's automatic differentiation support, and to be suitable for just-in-time compilation.	required
sample_initial_state	InitialStateSampler	Callable with one argument, which when passed a NumPy random number generator returns a random initial state for optimization of appropriate dimension.	required
rng	Generator	Seeded NumPy random number generator.	required

Other Parameters:

Name	Type	Description
number_minima_to_find	int	Number of candidate minima of objective function to try to find.
maximum_minimize_calls	int	Maximum number of times to try calling scipy.optimize.minimize to find candidate minima. If insufficient candidates are found within this number of calls then a ConvergenceError exception is raised.
minimize_method	str	String specifying one of local optimization methods which can be passed to method argument of scipy.optimize.minimize.
minimize_max_iterations	int | None	Maximum number of iterations in inner local minimization.
minimize_tol	float | None	Tolerance parameter for inner local minimization.
logging_function	Callable[[str], None]	Function to use to optionally log status messages during minimization. Defaults to a no-op function which discards messages.

Returns:

Type	Description
Array	Tuple with first entry the state corresponding to the best minima candidate
float	found and the second entry the corresponding objective function value.

Source code in src/calibr/optimization.py

def minimize_with_restarts(
    objective_function: ObjectiveFunction,
    sample_initial_state: InitialStateSampler,
    rng: Generator,
    *,
    number_minima_to_find: int = 5,
    maximum_minimize_calls: int = 100,
    minimize_method: str = "Newton-CG",
    minimize_max_iterations: int | None = None,
    minimize_tol: float | None = None,
    logging_function: Callable[[str], None] = lambda _: None,
    **unknown_kwargs: GlobalMinimizerKwarg,
) -> tuple[jax.Array, float]:
    """Minimize a differentiable objective function with random restarts.

    Iteratively calls `scipy.optimize.minimize` to attempt to find a minimum of an
    objective function until a specified number of candidate minima are successfully
    found, with the initial state for each `minimize` called being randomly sampled
    using a user provided function. The candidate minima with the minimum value for
    the objective function is returned along with the corresponding objective function
    value.

    Args:
        objective_function: Differentiable scalar-valued function of a single flat
            vector argument to be minimized. Assumed to be specified using JAX
            primitives such that its gradient and Hessian can be computed using JAX's
            automatic differentiation support, and to be suitable for just-in-time
            compilation.
        sample_initial_state: Callable with one argument, which when passed a NumPy
            random number generator returns a random initial state for optimization of
            appropriate dimension.
        rng: Seeded NumPy random number generator.

    Keyword Args:
        number_minima_to_find: Number of candidate minima of objective function to try
            to find.
        maximum_minimize_calls: Maximum number of times to try calling
            `scipy.optimize.minimize` to find candidate minima. If insufficient
            candidates are found within this number of calls then a `ConvergenceError`
            exception is raised.
        minimize_method: String specifying one of local optimization methods which can
            be passed to `method` argument of `scipy.optimize.minimize`.
        minimize_max_iterations: Maximum number of iterations in inner local
            minimization.
        minimize_tol: Tolerance parameter for inner local minimization.
        logging_function: Function to use to optionally log status messages during
            minimization. Defaults to a no-op function which discards messages.

    Returns:
        Tuple with first entry the state corresponding to the best minima candidate
        found and the second entry the corresponding objective function value.
    """
    _check_unknown_kwargs(unknown_kwargs)
    minima_found: list[tuple[jax.Array, int, jax.Array]] = []
    minimize_calls = 0
    while (
        len(minima_found) < number_minima_to_find
        and minimize_calls < maximum_minimize_calls
    ):
        logging_function(f"Starting minimize call {minimize_calls + 1}")
        results = _minimize(
            jax.jit(objective_function),
            x0=sample_initial_state(rng),
            jac=jax.jit(jax.grad(objective_function)),
            hessp=jax.jit(hessian_vector_product(objective_function)),
            method=minimize_method,
            tol=minimize_tol,
            options={"maxiter": minimize_max_iterations},
        )
        minimize_calls += 1
        if results.success:
            logging_function(f"Found minima with value {results.fun}")
            # Add minima to minima_found maintaining heap invariant such that first
            # entry in minima_found is always best solution so far (with counter used
            # to break ties between solutions with equal values for objective)
            heappush(minima_found, (results.fun, len(minima_found), results.x))
        else:
            logging_function(f"Minimization unsuccessful - {results.message}")
    if len(minima_found) < number_minima_to_find:
        msg = (
            f"Did not find required {number_minima_to_find} minima in "
            f"{maximum_minimize_calls} minimize calls."
        )
        raise ConvergenceError(msg)
    # Heap property means first entry in minima_found will correspond to solution
    # with minimum acquisition function value
    (
        min_objective_function,
        _,
        state,
    ) = minima_found[0]
    logging_function(f"Best minima found has value {min_objective_function}")
    return state, float(min_objective_function)