def calc_expected_metric_tensor_single(x):
if len(x.shape) == 1:
x = x.reshape(1, input_dim)
mu_j, cov_j = gp_jacobian(
# Xnew,
x,
X,
kernel,
mean_func,
f=f,
full_cov=full_cov,
# full_cov=False,
q_sqrt=q_sqrt,
jitter=jitter,
white=True,
)
# mu_j, cov_j = gp_jacobian_hard_coded(cov_fn, x, X, Y, jitter=jitter)
assert mu_j.shape == (input_dim, 1)
assert cov_j.shape == (input_dim, input_dim)
expected_jac_outer_prod = np.matmul(mu_j,
mu_j.T) # [input_dim x input_dim]
assert expected_jac_outer_prod.shape == (input_dim, input_dim)
# expected_metric_tensor = expected_jac_outer_prod + cov_weight * output_dim * cov_j
expected_metric_tensor = expected_jac_outer_prod + cov_weight * cov_j
# expected_metric_tensor = cov_weight * cov_j.T
# expected_metric_tensor = cov_j
assert expected_metric_tensor.shape == (input_dim, input_dim)
return expected_metric_tensor, mu_j, cov_j
def gp_jacobian_all(x):
if len(x.shape) == 1:
x = x.reshape(1, -1)
return gp_jacobian(
x,
gp.Z,
gp.kernel,
gp.mean_func,
f=gp.q_mu,
q_sqrt=gp.q_sqrt,
full_cov=False,
)