Skip to content

Commit

Permalink
Merge pull request #354 from mir-group/feature/yu/varmap_qr
Browse files Browse the repository at this point in the history 
QR for numerical stability of mapped variance
  • Loading branch information
@jonpvandermause
jonpvandermause authored Sep 16, 2024
2 parents 275c7ab + f0c09fe commit 48d9cdd
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 129 deletions.
33 changes: 18 additions & 15 deletions lammps_plugins/compute_flare_std_atom.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,17 @@ void ComputeFlareStdAtom::compute_peratom() {
continue;

if (use_map) {
int power = 2;
compute_energy_and_u(B2_vals, B2_norm_squared, single_bond_vals, power,
n_species, n_max, l_max, beta_matrices[itype - 1], u, &variance, normalized);
variance /= sig2;
Eigen::VectorXd Q_desc;
double K_self;
if (normalized) {
K_self = 1.0;
double B2_norm = pow(B2_norm_squared, 0.5);
Q_desc = beta_matrices[itype - 1].transpose() * B2_vals / B2_norm;
} else {
K_self = B2_norm_squared; // only power 1 is supported
Q_desc = beta_matrices[itype - 1].transpose() * B2_vals;
}
variance = K_self - Q_desc.dot(Q_desc);
} else {
Eigen::VectorXd kernel_vec = Eigen::VectorXd::Zero(n_clusters);
double K_self;
Expand Down Expand Up @@ -462,18 +469,14 @@ void ComputeFlareStdAtom::read_file(char *filename) {
int beta_count = 0;
double beta_val;
for (int k = 0; k < n_species; k++) {
// for (int l = 0; l < n_species; l++) {

beta_matrix = Eigen::MatrixXd::Zero(n_descriptors, n_descriptors);
for (int i = 0; i < n_descriptors; i++) {
for (int j = 0; j < n_descriptors; j++) {
//beta_matrix(k * n_descriptors + i, l * n_descriptors + j) = beta[beta_count];
beta_matrix(i, j) = beta[beta_count];
beta_count++;
}
beta_matrix = Eigen::MatrixXd::Zero(n_descriptors, n_descriptors);
for (int i = 0; i < n_descriptors; i++) {
for (int j = 0; j < n_descriptors; j++) {
beta_matrix(i, j) = beta[beta_count];
beta_count++;
}
beta_matrices.push_back(beta_matrix);
// }
}
beta_matrices.push_back(beta_matrix);
}

}
Expand Down
65 changes: 12 additions & 53 deletions src/flare_pp/kernels/dot_product.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -786,8 +786,8 @@ Eigen::MatrixXd DotProduct ::compute_varmap_coefficients(
int beta_size = p_size * (p_size + 1) / 2;
int n_species = gp_model.sparse_descriptors[kernel_index].n_types;
int n_sparse = gp_model.sparse_descriptors[kernel_index].n_clusters;
//mapping_coeffs = Eigen::MatrixXd::Zero(n_species * n_species, p_size * p_size); // can be reduced by symmetry
mapping_coeffs = Eigen::MatrixXd::Zero(n_species, p_size * p_size); // can be reduced by symmetry
double jitter_root = pow(gp_model.Kuu_jitter, 0.5);

// Get alpha index.

Expand All @@ -799,58 +799,17 @@ Eigen::MatrixXd DotProduct ::compute_varmap_coefficients(
// Loop over types.
for (int s = 0; s < n_species; s++){
int n_types = gp_model.sparse_descriptors[kernel_index].n_clusters_by_type[s];
int c_types =
gp_model.sparse_descriptors[kernel_index].cumulative_type_count[s];
int K_ind = alpha_ind + c_types;

// Loop over clusters within each type.
for (int i = 0; i < n_types; i++){
Eigen::VectorXd pi_current =
gp_model.sparse_descriptors[kernel_index].descriptors[s].row(i);
double pi_norm =
gp_model.sparse_descriptors[kernel_index].descriptor_norms[s](i);

// Skip empty environments.
if (pi_norm < empty_thresh)
continue;

// TODO: include symmetry of i & j
// Loop over clusters within each type.
for (int j = 0; j < n_types; j++){
Eigen::VectorXd pj_current =
gp_model.sparse_descriptors[kernel_index].descriptors[s].row(j);
double pj_norm =
gp_model.sparse_descriptors[kernel_index].descriptor_norms[s](j);

// Skip empty environments.
if (pj_norm < empty_thresh)
continue;

double Kuu_inv_ij = gp_model.Kuu_inverse(K_ind + i, K_ind + j);
double Kuu_inv_ij_normed = Kuu_inv_ij; // / pi_norm / pj_norm;
// double Sigma_ij = gp_model.Sigma(K_ind + i, K_ind + j);
// double Sigma_ij_normed = Sigma_ij / pi_norm / pj_norm;
int beta_count = 0;

// First loop over descriptor values.
for (int k = 0; k < p_size; k++) {
double p_ik = pi_current(k);

// Second loop over descriptor values.
for (int l = 0; l < p_size; l++){
double p_jl = pj_current(l);

// Update beta vector.
double beta_val = sig2 * sig2 * p_ik * p_jl * (- Kuu_inv_ij_normed); // + Sigma_ij_normed); // To match the compute_cluster_uncertainty function
mapping_coeffs(s, beta_count) += beta_val;

if (k == l && i == 0 && j == 0) {
mapping_coeffs(s, beta_count) += sig2; // the self kernel term
}

beta_count++;
}
}
Eigen::MatrixXd pi = gp_model.sparse_descriptors[kernel_index].descriptors[s];
Eigen::MatrixXd A = Eigen::MatrixXd::Zero(p_size + n_types, n_types);
A.block(0, 0, p_size, n_types) = pi.transpose();
A.block(p_size, 0, n_types, n_types) = jitter_root * Eigen::MatrixXd::Identity(n_types, n_types) / sigma;
Eigen::HouseholderQR<Eigen::MatrixXd> qr(A);
Eigen::MatrixXd Q = qr.householderQ(); // (p_size + n_types, n_types)
int beta_count = 0;
for (int i = 0; i < p_size; i++) {
for (int j = 0; j < p_size; j++) {
if (j < n_types) mapping_coeffs(s, beta_count) += Q(i, j);
beta_count++;
}
}
}
Expand Down
58 changes: 14 additions & 44 deletions src/flare_pp/kernels/normalized_dot_product.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -784,8 +784,8 @@ Eigen::MatrixXd NormalizedDotProduct ::compute_varmap_coefficients(
int beta_size = p_size * (p_size + 1) / 2;
int n_species = gp_model.sparse_descriptors[kernel_index].n_types;
int n_sparse = gp_model.sparse_descriptors[kernel_index].n_clusters;
//mapping_coeffs = Eigen::MatrixXd::Zero(n_species * n_species, p_size * p_size); // can be reduced by symmetry
mapping_coeffs = Eigen::MatrixXd::Zero(n_species, p_size * p_size); // can be reduced by symmetry
double jitter_root = pow(gp_model.Kuu_jitter, 0.5);

// Get alpha index.

Expand All @@ -797,58 +797,28 @@ Eigen::MatrixXd NormalizedDotProduct ::compute_varmap_coefficients(
// Loop over types.
for (int s = 0; s < n_species; s++){
int n_types = gp_model.sparse_descriptors[kernel_index].n_clusters_by_type[s];
int c_types =
gp_model.sparse_descriptors[kernel_index].cumulative_type_count[s];
int K_ind = alpha_ind + c_types;

Eigen::MatrixXd pi = gp_model.sparse_descriptors[kernel_index].descriptors[s];
// Loop over clusters within each type.
for (int i = 0; i < n_types; i++){
Eigen::VectorXd pi_current =
gp_model.sparse_descriptors[kernel_index].descriptors[s].row(i);
double pi_norm =
gp_model.sparse_descriptors[kernel_index].descriptor_norms[s](i);

// Skip empty environments.
if (pi_norm < empty_thresh)
continue;
pi.row(i) /= pi_norm; // normalize descriptor by norm
}

// TODO: include symmetry of i & j
// Loop over clusters within each type.
for (int j = 0; j < n_types; j++){
Eigen::VectorXd pj_current =
gp_model.sparse_descriptors[kernel_index].descriptors[s].row(j);
double pj_norm =
gp_model.sparse_descriptors[kernel_index].descriptor_norms[s](j);

// Skip empty environments.
if (pj_norm < empty_thresh)
continue;

double Kuu_inv_ij = gp_model.Kuu_inverse(K_ind + i, K_ind + j);
double Kuu_inv_ij_normed = Kuu_inv_ij / pi_norm / pj_norm;
// double Sigma_ij = gp_model.Sigma(K_ind + i, K_ind + j);
// double Sigma_ij_normed = Sigma_ij / pi_norm / pj_norm;
int beta_count = 0;

// First loop over descriptor values.
for (int k = 0; k < p_size; k++) {
double p_ik = pi_current(k);

// Second loop over descriptor values.
for (int l = 0; l < p_size; l++){
double p_jl = pj_current(l);

// Update beta vector.
double beta_val = sig2 * sig2 * p_ik * p_jl * (- Kuu_inv_ij_normed); // + Sigma_ij_normed); // To match the compute_cluster_uncertainty function
mapping_coeffs(s, beta_count) += beta_val;

if (k == l && i == 0 && j == 0) {
mapping_coeffs(s, beta_count) += sig2; // the self kernel term
}

beta_count++;
}
}
Eigen::MatrixXd A = Eigen::MatrixXd::Zero(p_size + n_types, n_types);
A.block(0, 0, p_size, n_types) = pi.transpose();
A.block(p_size, 0, n_types, n_types) = jitter_root * Eigen::MatrixXd::Identity(n_types, n_types) / sigma;
Eigen::HouseholderQR<Eigen::MatrixXd> qr(A);
Eigen::MatrixXd Q = qr.householderQ(); // (p_size + n_types, n_types)
int beta_count = 0;
for (int i = 0; i < p_size; i++) {
for (int j = 0; j < p_size; j++) {
if (j < n_types) mapping_coeffs(s, beta_count) += Q(i, j);
beta_count++;
}
}
}
Expand Down
31 changes: 20 additions & 11 deletions tests/get_sgp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,16 @@
from ase.build import make_supercell

# Define kernel.
sigma = 2.0
sigma = np.random.rand() + 1.0
power = 1.0
dotprod_kernel = DotProduct(sigma, power)
normdotprod_kernel = NormalizedDotProduct(sigma, power)

# Define remaining parameters for the SGP wrapper.
sigma_e = 0.3
sigma_f = 0.2
sigma_s = 0.1
species_map = {6: 0, 8: 1}
single_atom_energies = {0: -5, 1: -6}
variance_type = "local"
max_iterations = 20
opt_method = "L-BFGS-B"
bounds = [(None, None), (sigma_e, None), (None, None), (None, None)]
bounds = [(None, None), (None, None), (None, None), (None, None)]


def get_random_atoms(a=2.0, sc_size=2, numbers=[6, 8], set_seed: int = None):
Expand Down Expand Up @@ -89,6 +84,18 @@ def get_empty_sgp(n_types=2, power=2, multiple_cutoff=False, kernel_type="Normal
cutoff_matrix,
)

sigma_e = np.random.rand()
sigma_f = np.random.rand()
sigma_s = np.random.rand()
print("Hyps", kernel.sigma, sigma_e, sigma_f, sigma_s)

if n_types == 1:
species_map = {6: 0}
single_atom_energies = {0: np.random.rand()}
elif n_types == 2:
species_map = {6: 0, 8: 1}
single_atom_energies = {0: np.random.rand(), 1: np.random.rand()}

empty_sgp = SGP_Wrapper(
[kernel],
[b2_calc],
Expand Down Expand Up @@ -123,16 +130,18 @@ def get_updated_sgp(n_types=2, power=2, multiple_cutoff=False, kernel_type="Norm
energy = training_structure.get_potential_energy()
stress = training_structure.get_stress()

size = max(1, np.random.randint(len(training_structure)))
custom_range = np.random.choice(len(training_structure), size=size, replace=False).tolist()
sgp.update_db(
training_structure,
forces,
custom_range=(1, 2, 3, 4, 5),
custom_range=custom_range,
energy=energy,
stress=stress,
mode="specific",
rel_e_noise=0.1,
rel_f_noise=0.2,
rel_s_noise=0.1,
rel_e_noise=np.random.rand(),
rel_f_noise=np.random.rand(),
rel_s_noise=np.random.rand(),
)

# add an isolated atom to the training data
Expand Down
12 changes: 6 additions & 6 deletions tests/test_lammps.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@
@pytest.mark.parametrize("struc", struc_list)
@pytest.mark.parametrize("multicut", [False, True])
@pytest.mark.parametrize("n_cpus", n_cpus_list)
def test_write_potential(n_species, n_types, power, struc, multicut, n_cpus):
@pytest.mark.parametrize("kernel_type", ["NormalizedDotProduct", "DotProduct"])
def test_write_potential(n_species, n_types, power, struc, multicut, n_cpus, kernel_type):
"""Test the flare pair style."""

if n_species > n_types:
Expand Down Expand Up @@ -100,11 +101,11 @@ def test_write_potential(n_species, n_types, power, struc, multicut, n_cpus):
energy_lmp += sgp_model.gp_model.single_atom_energies[coded_spec]

thresh = 1e-6
r_thresh = 1e-3
print(energy, energy_lmp)
assert np.allclose(energy, energy_lmp, atol=thresh)
# print(forces, forces_lmp)
assert np.allclose(forces, forces_lmp, atol=thresh)
assert np.allclose(stress, stress_lmp, atol=thresh)
assert np.allclose(energy, energy_lmp, atol=thresh, rtol=r_thresh)
assert np.allclose(forces, forces_lmp, atol=thresh, rtol=r_thresh)
assert np.allclose(stress, stress_lmp, atol=thresh, rtol=r_thresh)

# Remove files.
os.remove(potential_name)
Expand All @@ -119,7 +120,6 @@ def test_write_potential(n_species, n_types, power, struc, multicut, n_cpus):
from ase.calculators.lammpsrun import LAMMPS
from flare.md.lammps import LAMMPS_MOD, LAMMPS_MD, get_kinetic_stress


@pytest.mark.skipif(
not os.environ.get("lmp", False),
reason=(
Expand Down

0 comments on commit 48d9cdd

Please sign in to comment.