38 |
39 |
Source code for bayesquad.acquisition_functions
40 | """Acquisition functions and related functions."""
41 | import numpy as np
42 |
43 | from .plotting import plottable
44 | from .quadrature import IntegrandModel
45 |
46 |
47 | [docs]def model_variance(integrand_model: IntegrandModel):
48 |
49 |
@plottable("Model variance", default_plotting_parameters={'calculate_jacobian': False})
50 |
def f(x, *, calculate_jacobian=True):
51 |
"""Evaluate the variance, and the Jacobian of the variance, for the given `IntegrandModel` at a point, or a set
52 |
of points.
53 |
54 |
Given an array of shape (num_points, num_dimensions), returns an array of shape (num_points) containing the
55 |
function values and an array of shape (num_points, num_dimensions) containing the function Jacobians.
56 |
57 |
Given an array of shape (num_dimensions), returns a 0D array containing the function value and an array of shape
58 |
(num_dimensions) containing the function Jacobian.
59 |
60 |
If the Jacobian is not required (e.g. for plotting), the relevant calculations can be disabled by setting
61 |
`calculate_jacobian=False`.
62 |
"""
63 |
_, variance = integrand_model.posterior_mean_and_variance(x)
64 |
65 |
if calculate_jacobian:
66 |
variance_jacobian = integrand_model.posterior_variance_jacobian(x)
67 |
else:
68 |
variance_jacobian = None
69 |
70 |
return variance, variance_jacobian
71 |
72 |
return f
73 |
74 |
75 | [docs]def model_variance_norm_of_gradient_squared(integrand_model: IntegrandModel):
76 |
77 |
@plottable("Gradient squared", default_plotting_parameters={'calculate_jacobian': False})
78 |
def f(x, *, calculate_jacobian=True):
79 |
"""Evaluate the squared norm of the gradient of the variance, and the Jacobian of this quantity, for the given
80 |
`IntegrandModel` at a point, or a set of points.
81 |
82 |
Given an array of shape (num_points, num_dimensions), returns an array of shape (num_points) containing the
83 |
function values and an array of shape (num_points, num_dimensions) containing the function Jacobians.
84 |
85 |
Given an array of shape (num_dimensions), returns a 0D array containing the function value and an array of shape
86 |
(num_dimensions) containing the function Jacobian.
87 |
88 |
If the Jacobian is not required (e.g. for plotting), the relevant calculations can be disabled by setting
89 |
`calculate_jacobian=False`.
90 |
"""
91 |
variance_jacobian = integrand_model.posterior_variance_jacobian(x)
92 |
93 |
# Inner product of the Jacobian with itself, for each point.
94 |
gradient_squared = np.einsum('...i,...i->...', variance_jacobian, variance_jacobian, optimize=True)
95 |
96 |
if calculate_jacobian:
97 |
variance_hessian = integrand_model.posterior_variance_hessian(x)
98 |
99 |
# Matrix product of Hessian and Jacobian, for each point.
100 |
gradient_squared_jacobian = 2 * np.einsum('...ij,...j->...i',
101 |
variance_hessian,
102 |
variance_jacobian,
103 |
optimize=True)
104 |
else:
105 |
gradient_squared_jacobian = None
106 |
107 |
return gradient_squared, gradient_squared_jacobian
108 |
109 |
return f
110 |
111 |
112 |
