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 byPyMC
andNumPyro
to deal with parameters bounded by a lower and upper bound.Example
param_bounds
should be of the formparam_bounds[par] = [lower_bound, upper_bound]
wherelower_bound
andupper_bound
are of the same shape asp[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
ofstr
, optional) – Thelist
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
ofstr
, 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 anException
.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 toTrue
.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]
andcov_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 whenreturn_array = False
for consistency. The order of the list matches howjax.flatten_util.ravel_pytree
will order keys from a PyTree.- Return type:
(JAXArray,
list
ofstr
) or (PyTree,list
ofstr
)