def test_ky_mat(params, ihyps, ky_mat_ref):
name, cutoffs, hyps_mask_list, energy_noise = params
hyps_mask = hyps_mask_list[ihyps]
hyps = hyps_mask["hyps"]
kernel = str_to_kernel_set(hyps_mask["kernels"], "mc", hyps_mask)
# serial implementation
time0 = time.time()
ky_mat = get_Ky_mat(hyps, name, kernel[0], kernel[2], kernel[3],
energy_noise, cutoffs, hyps_mask)
print(f"compute ky_mat with multihyps, test {ihyps}, n_cpus=1",
time.time() - time0)
assert np.isclose(ky_mat, ky_mat_ref, rtol=1e-3).all(
), f"multi hyps implementation is wrongwith case {ihyps}"
# parallel implementation
time0 = time.time()
ky_mat = get_Ky_mat(
hyps,
name,
kernel[0],
kernel[2],
kernel[3],
energy_noise,
cutoffs,
hyps_mask,
n_cpus=2,
n_sample=20,
)
print(f"compute ky_mat with multihyps, test {ihyps}, n_cpus=2",
time.time() - time0)
assert np.isclose(ky_mat, ky_mat_ref, rtol=1e-3).all(
), f"multi hyps parallel implementation is wrong with case {ihyps}"
def ky_mat_ref(params):
name, cutoffs, hyps_mask_list, energy_noise = params
# get the reference without multi hyps
hyps_mask = hyps_mask_list[-1]
hyps = hyps_mask["hyps"]
kernel = str_to_kernel_set(hyps_mask["kernels"], "mc", hyps_mask)
time0 = time.time()
ky_mat0 = get_Ky_mat(hyps, name, kernel[0], kernel[2], kernel[3],
energy_noise, cutoffs)
print("compute ky_mat serial", time.time() - time0)
# parallel version
time0 = time.time()
ky_mat = get_Ky_mat(
hyps,
name,
kernel[0],
kernel[2],
kernel[3],
energy_noise,
cutoffs,
n_cpus=2,
n_sample=5,
)
print("compute ky_mat parallel", time.time() - time0)
assert np.isclose(ky_mat, ky_mat0,
rtol=1e-3).all(), "parallel implementation is wrong"
yield ky_mat0
del ky_mat0
def set_L_alpha(self):
"""
Invert the covariance matrix, setting L (a lower triangular
matrix s.t. L L^T = (K + sig_n^2 I)) and alpha, the inverse
covariance matrix multiplied by the vector of training labels.
The forces and variances are later obtained using alpha.
"""
self.sync_data()
ky_mat = get_Ky_mat(
self.hyps,
self.name,
self.kernel,
self.energy_kernel,
self.energy_force_kernel,
self.energy_noise,
cutoffs=self.cutoffs,
hyps_mask=self.hyps_mask,
n_cpus=self.n_cpus,
n_sample=self.n_sample,
)
l_mat = np.linalg.cholesky(ky_mat)
l_mat_inv = np.linalg.inv(l_mat)
ky_mat_inv = l_mat_inv.T @ l_mat_inv
alpha = np.matmul(ky_mat_inv, self.all_labels)
self.ky_mat = ky_mat
self.l_mat = l_mat
self.alpha = alpha
self.ky_mat_inv = ky_mat_inv
self.likelihood = get_like_from_mats(ky_mat, l_mat, alpha, self.name)
self.n_envs_prev = len(self.training_data)
def set_L_alpha_part(self, expert_id):
"""
Invert the covariance matrix, setting L (a lower triangular
matrix s.t. L L^T = (K + sig_n^2 I)) and alpha, the inverse
covariance matrix multiplied by the vector of training labels.
The forces and variances are later obtained using alpha.
"""
self.sync_experts_data(expert_id)
if self.per_expert_parallel and self.n_cpus > 1:
n_cpus = 1
else:
n_cpus = self.n_cpus
ky_mat = get_Ky_mat(
self.hyps,
f"{self.name}_{expert_id}",
self.kernel,
self.energy_kernel,
self.energy_force_kernel,
self.energy_noise,
cutoffs=self.cutoffs,
hyps_mask=self.hyps_mask,
n_cpus=n_cpus,
n_sample=self.n_sample,
)
self.compute_experts_matrices(ky_mat, expert_id)
self.likelihood[expert_id] = get_like_from_mats(
self.ky_mat[expert_id],
self.l_mat[expert_id],
self.alpha[expert_id],
f"{self.name}_{expert_id}",
)
def set_L_alpha(self):
"""
Set the lower-triangular (L) version of the covariance matrix and alpha vector (
mapping from similarity vector of a new training point to output label) for
each expert.
:return:
"""
self.sync_data()
logger = logging.getLogger(self.logger_name)
logger.debug("set_L_alpha")
if self.per_expert_parallel and self.n_cpus > 1:
time0 = time.time()
with mp.Pool(processes=self.n_cpus) as pool:
results = []
for expert_id in range(self.n_experts):
results.append(
pool.apply_async(
get_Ky_mat,
(
self.hyps,
f"{self.name}_{expert_id}",
self.kernel,
self.energy_kernel,
self.energy_force_kernel,
self.energy_noise,
self.cutoffs,
self.hyps_mask,
1,
self.n_sample,
),
))
for i in range(self.n_experts):
ky_mat = results[i].get()
self.compute_experts_matrices(ky_mat, i)
pool.close()
pool.join()
logger.debug(
f"set_L_alpha with per_expert_par {time.time()-time0}")
else:
for expert_id in range(self.n_experts):
logger.debug(f"compute L_alpha for {expert_id}")
time0 = time.time()
ky_mat = get_Ky_mat(
self.hyps,
f"{self.name}_{expert_id}",
self.kernel,
self.energy_kernel,
self.energy_force_kernel,
self.energy_noise,
self.cutoffs,
self.hyps_mask,
self.n_cpus,
self.n_sample,
)
self.compute_experts_matrices(ky_mat, expert_id)
logger.debug(f"{expert_id} compute_L_alpha {time.time()-time0}"
f" {len(self.training_data[expert_id])}")
for expert_id in range(self.n_experts):
self.likelihood[expert_id] = get_like_from_mats(
self.ky_mat[expert_id],
self.l_mat[expert_id],
self.alpha[expert_id],
f"{self.name}_{expert_id}",
)
self.total_likelihood = np.sum(self.likelihood)
