GP.laplace_approx_with_bounds

GP.laplace_approx_with_bounds#

GP.laplace_approx_with_bounds(p: Any, Y: Array, param_bounds: Any, vars: list | None = None, fixed_vars: list | None = None, large: bool | None = False, large_block_size: int | None = 50, return_array: bool | None = False, large_jit: bool | None = True, **kwargs) → Tuple[Any | Array, list][source]#

Computes the Laplace approximation at the location of p but within the transformed parameter space used by PyMC and NumPyro to deal with parameters bounded by a lower and upper bound.

Example

param_bounds should be of the form param_bounds[par] = [lower_bound, upper_bound] where lower_bound and upper_bound are of the same shape as p[par].

See GP.laplace_approx for more details about the Laplace approximation.

Parameters:

  • p (PyTree) – Pytree of hyperparameters used to calculate the covariance matrix in addition to any mean function parameters which may be needed to calculate the mean function. Also input to the logPrior function for the calculation of the log priors.

  • Y (JAXArray) – Observed data to fit, must be of shape (N_l, N_t).

  • param_bounds (PyTree) – Contains any bounds for the parameters in p.

  • vars (list of str, optional) – The list of key names corresponding to the parameters we want to calculate the Laplace approximation with respect to. The remaining parameters will be assumed to be fixed. If specified in addition to fixed_vars will raise an Exception.

  • fixed_vars (list of str, optional) – Alternative to vars, may specify instead the parameters being kept fixed which will not be marginalised over in the Laplace approximation. If specified in addition to vars will raise an Exception.

  • large (bool, optional) – Calculating the hessian matrix for large data sets with many parameters can be very memory intensive. If this is set to True then the hessian will be calculated in groups of rows instead of all at once which reduces the memory cost but can take significantly longer to run. The calculation is otherwise the same with no approximation made. Defaults to False.

  • large_block_size (int, optional) – If large is set to True and the hessian is being calculated in groups of rows can specify how many rows are being calculated simultaneously. Large numbers may calculate the overall hessian faster but at greater memory cost.

  • large_jit (bool, optional) – Whether to JIT compile the hessian function when large = True, can speed up the calculation assuming the function can be JIT compiled. Defaults to True.

  • return_array (bool, optional) – Whether to return the approximated covariance matrix as a JAXArray or as a nested PyTree where e.g. the covariance between parameters named p1 and p2 is given by cov_mat[p1][p2] and cov_mat[p2][p1].

Returns:

Returns a tuple of two elements, if return_array = True the first element will be the covariance matrix from the Laplace approximation as a JAXArray, otherwise it will be as a nested PyTree. The second element will be the order of the parameters in the returned covariance matrix if it is a JAXArray. This list is also returned when return_array = False for consistency. The order of the list matches how jax.flatten_util.ravel_pytree will order keys from a PyTree.

Return type:

(JAXArray, list of str) or (PyTree, list of str)

Contents