Skip to content

Commit

Permalink
Merge pull request #166 from mir-group/bugfix/mb-etypes
Browse files Browse the repository at this point in the history 
fix manybody kernel etype issue
  • Loading branch information
@nw13slx
nw13slx authored Apr 30, 2020
2 parents 27e8c4c + 7bdeb6d commit f907c9a
Show file tree
Hide file tree
Showing 8 changed files with 286 additions and 1,459 deletions.
7 changes: 4 additions & 3 deletions flare/env.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,8 @@ def compute_env(self):

# if 3 cutoffs are given, create many-body arrays
if len(self.cutoffs) > 2:
self.bond_array_mb, self.neigh_dists_mb, self.num_neighs_mb, self.etype_mb = get_m_body_arrays(
self.bond_array_mb, self.neigh_dists_mb, self.num_neighs_mb, self.etype_mb, \
self.bond_array_mb_etypes = get_m_body_arrays(
self.positions, self.atom, self.cell, self.cutoffs[2], self.species)
else:
self.bond_array_mb = None
Expand Down Expand Up @@ -412,7 +413,7 @@ class to allow for njit acceleration with Numba.
"""
# TODO: this can be probably improved using stored arrays, redundant calls to get_2_body_arrays
# Get distances, positions, species and indexes of neighbouring atoms
bond_array_mb, __, _, bond_inds = get_2_body_arrays_ind(
bond_array_mb, __, etypes, bond_inds = get_2_body_arrays_ind(
positions, atom, cell, cutoff_mb, species)

# For each neighbouring atom, get distances in its neighbourhood
Expand All @@ -437,7 +438,7 @@ class to allow for njit acceleration with Numba.
etypes_mb_array[i, :num_neighs_mb[i]] = neighbouring_etypes[i]


return bond_array_mb, neigh_dists_mb, num_neighs_mb, etypes_mb_array
return bond_array_mb, neigh_dists_mb, num_neighs_mb, etypes_mb_array, etypes


if __name__ == '__main__':
Expand Down
37 changes: 22 additions & 15 deletions flare/kernels/mc_sephyps.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,10 +94,13 @@ def two_three_many_body_mc(env1, env2, d1, d2, cutoffs,
d1, d2, sig3, ls3, r_cut_3, cutoff_func,
nspec, spec_mask, triplet_mask)

many_term = many_body_mc_jit(env1.bond_array_mb, env2.bond_array_mb, env1.neigh_dists_mb,
env2.neigh_dists_mb, env1.num_neighs_mb, env2.num_neighs_mb,
env1.ctype, env2.ctype, env1.etypes, env2.etypes,
env1.etype_mb, env2.etype_mb, env1.species, env2.species,
many_term = many_body_mc_jit(env1.bond_array_mb, env2.bond_array_mb,
env1.neigh_dists_mb, env2.neigh_dists_mb,
env1.num_neighs_mb, env2.num_neighs_mb,
env1.ctype, env2.ctype,
env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.etype_mb, env2.etype_mb,
env1.species, env2.species,
d1, d2, sigm, lsm, r_cut_m, cutoff_func)

return two_term + three_term + many_term
Expand Down Expand Up @@ -133,11 +136,12 @@ def two_three_many_body_mc_grad(env1, env2, d1, d2, cutoffs,

kern_many, gradm = many_body_mc_grad_jit(env1.bond_array_mb, env2.bond_array_mb,
env1.neigh_dists_mb, env2.neigh_dists_mb,
env1.num_neighs_mb, env2.num_neighs_mb, env1.ctype,
env2.ctype, env1.etypes, env2.etypes,
env1.num_neighs_mb, env2.num_neighs_mb,
env1.ctype, env2.ctype,
env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.etype_mb, env2.etype_mb,
env1.species, env2.species, d1, d2, sigm,
lsm, r_cut_m, cutoff_func)
env1.species, env2.species,
d1, d2, sigm, lsm, r_cut_m, cutoff_func)

g = np.hstack([grad2, grad3, gradm])

Expand Down Expand Up @@ -173,11 +177,13 @@ def two_three_many_mc_force_en(env1, env2, d1, cutoffs,
triplet_mask) / 3

many_term = many_body_mc_force_en_jit(env1.bond_array_mb, env2.bond_array_mb,
env1.neigh_dists_mb, env1.num_neighs_mb,
env1.ctype, env2.ctype, env1.etypes, env2.etypes,
env1.neigh_dists_mb,
env1.num_neighs_mb,
env1.ctype, env2.ctype,
env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.etype_mb,
env1.species, env2.species, d1, sigm, lsm, r_cut_m,
cutoff_func)
env1.species, env2.species,
d1, sigm, lsm, r_cut_m, cutoff_func)

return two_term + three_term + many_term

Expand Down Expand Up @@ -208,9 +214,10 @@ def two_three_many_mc_en(env1, env2, cutoffs,
nspec, spec_mask,
triplet_mask)

many_term = many_body_mc_en_jit(env1.bond_array_2, env2.bond_array_2, env1.ctype,
env2.ctype, env1.etypes, env2.etypes, env1.species,
env2.species,
many_term = many_body_mc_en_jit(env1.bond_array_2, env2.bond_array_2,
env1.ctype, env2.ctype,
env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.species, env2.species,
sigm, lsm, r_cut_m, cutoff_func)


Expand Down
25 changes: 15 additions & 10 deletions flare/kernels/mc_simple.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,8 @@ def two_plus_three_plus_many_body_mc(env1: AtomicEnvironment, env2: AtomicEnviro

many_term = many_body_mc_jit(env1.bond_array_mb, env2.bond_array_mb, env1.neigh_dists_mb,
env2.neigh_dists_mb, env1.num_neighs_mb, env2.num_neighs_mb,
env1.ctype, env2.ctype, env1.etypes, env2.etypes,
env1.ctype, env2.ctype,
env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.etype_mb, env2.etype_mb, env1.species, env2.species,
d1, d2, sigm, lsm, r_cut_m, cutoff_func)

Expand Down Expand Up @@ -305,7 +306,8 @@ def two_plus_three_plus_many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicE
kern_many, gradm = many_body_mc_grad_jit(env1.bond_array_mb, env2.bond_array_mb,
env1.neigh_dists_mb, env2.neigh_dists_mb,
env1.num_neighs_mb, env2.num_neighs_mb, env1.ctype,
env2.ctype, env1.etypes, env2.etypes,
env2.ctype, env1.bond_array_mb_etypes,
env2.bond_array_mb_etypes,
env1.etype_mb, env2.etype_mb,
env1.species, env2.species, d1, d2, sigm,
lsm, r_cut_m, cutoff_func)
Expand Down Expand Up @@ -359,7 +361,8 @@ def two_plus_three_plus_many_body_mc_force_en(env1: AtomicEnvironment, env2: Ato

many_term = many_body_mc_force_en_jit(env1.bond_array_mb, env2.bond_array_mb,
env1.neigh_dists_mb, env1.num_neighs_mb,
env1.ctype, env2.ctype, env1.etypes, env2.etypes,
env1.ctype, env2.ctype, env1.bond_array_mb_etypes,
env2.bond_array_mb_etypes,
env1.etype_mb,
env1.species, env2.species, d1, sigm, lsm, r_cut_m,
cutoff_func)
Expand Down Expand Up @@ -407,8 +410,9 @@ def two_plus_three_plus_many_body_mc_en(env1: AtomicEnvironment, env2: AtomicEnv
env1.triplet_counts, env2.triplet_counts,
sig3, ls3, r_cut_3, cutoff_func)

many_term = many_body_mc_en_jit(env1.bond_array_2, env2.bond_array_2, env1.ctype,
env2.ctype, env1.etypes, env2.etypes, env1.species,
many_term = many_body_mc_en_jit(env1.bond_array_mb, env2.bond_array_mb, env1.ctype,
env2.ctype, env1.bond_array_mb_etypes,
env2.bond_array_mb_etypes, env1.species,
env2.species,
sigm, lsm, r_cut_m, cutoff_func)

Expand Down Expand Up @@ -707,7 +711,7 @@ def many_body_mc(env1: AtomicEnvironment, env2: AtomicEnvironment,

# Get atomic species of central atom, their neighbours, and their neighbours' neighbours
c1, c2 = env1.ctype, env2.ctype
etypes1, etypes2 = env1.etypes, env2.etypes
etypes1, etypes2 = env1.bond_array_mb_etypes, env2.bond_array_mb_etypes
etypes_neigh_1, etypes_neigh_2 = env1.etype_mb, env2.etype_mb

return many_body_mc_jit(bond_array_1, bond_array_2, neigh_dists_1, neigh_dists_2, num_neigh_1,
Expand Down Expand Up @@ -735,7 +739,7 @@ def many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicEnvironment,
num_neigh_2 = env2.num_neighs_mb

c1, c2 = env1.ctype, env2.ctype
etypes1, etypes2 = env1.etypes, env2.etypes
etypes1, etypes2 = env1.bond_array_mb_etypes, env2.bond_array_mb_etypes
etypes_neigh_1, etypes_neigh_2 = env1.etype_mb, env2.etype_mb

kernel, kernel_grad = many_body_mc_grad_jit(bond_array_1, bond_array_2, neigh_dists_1,
Expand Down Expand Up @@ -773,7 +777,7 @@ def many_body_mc_force_en(env1, env2, d1, hyps, cutoffs,
num_neigh_1 = env1.num_neighs_mb

c1, c2 = env1.ctype, env2.ctype
etypes1, etypes2 = env1.etypes, env2.etypes
etypes1, etypes2 = env1.bond_array_mb_etypes, env2.bond_array_mb_etypes
etypes_neigh_1 = env1.etype_mb

# divide by three to account for triple counting
Expand Down Expand Up @@ -802,8 +806,9 @@ def many_body_mc_en(env1: AtomicEnvironment, env2: AtomicEnvironment,
ls = hyps[1]
r_cut = cutoffs[2]

return many_body_mc_en_jit(env1.bond_array_2, env2.bond_array_2, env1.ctype,
env2.ctype, env1.etypes, env2.etypes, env1.species, env2.species,
return many_body_mc_en_jit(env1.bond_array_mb, env2.bond_array_mb, env1.ctype,
env2.ctype, env1.bond_array_mb_etypes, env2.bond_array_mb_etypes,
env1.species, env2.species,
sig, ls, r_cut, cutoff_func)


Expand Down
8 changes: 4 additions & 4 deletions flare/kernels/sc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1848,10 +1848,10 @@ def mb_grad_helper_ls(q1, q2, qi, qj, sig, ls):
'2+3_grad': two_plus_three_body_grad,
'2+3_en': two_plus_three_en,
'2+3_force_en': two_plus_three_force_en,
'many_body': many_body,
'many_body_en': many_body_en,
'many_body_grad': many_body_grad,
'many_body_force_en': many_body_en,
'many': many_body,
'many_en': many_body_en,
'many_grad': many_body_grad,
'many_force_en': many_body_force_en,
'two_plus_three_plus_many_body': two_plus_three_plus_many_body,
'two_plus_three_plus_many_body_grad': two_plus_three_plus_many_body_grad,
'two_plus_three_plus_many_body_en': two_plus_three_plus_many_body_en,
Expand Down
84 changes: 74 additions & 10 deletions tests/fake_gp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from flare.struc import Structure
from flare.otf_parser import OtfAnalysis


def get_random_structure(cell, unique_species, noa):
"""Create a random test structure """
np.random.seed(0)
Expand All @@ -35,11 +36,11 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru
if (nbond > 0):
nbond = 1
hyps_label += ['Length', 'Signal Var.']
if (ntriplet >0):
if (ntriplet > 0):
ntriplet = 1
hyps_label += ['Length', 'Signal Var.']
hyps_label += ['Noise Var.']
return random((nbond+ntriplet)*2+1), {'hyps_label':hyps_label}, deepcopy(cutoffs)
return random((nbond+ntriplet)*2+1), {'hyps_label': hyps_label}, deepcopy(cutoffs)

specs_mask = np.zeros(118, dtype=int)
specs_mask[1] = 0
Expand All @@ -55,7 +56,7 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru
cut += [cutoffs[0]]
cut += [cutoffs[1]]

if (nbond>1):
if (nbond > 1):
sig1 = random(nbond)
ls1 = random(nbond)
bond_mask = np.ones(nspecs**2, dtype=int)
Expand All @@ -67,7 +68,7 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru
bond_mask = np.zeros(nspecs**2, dtype=int)
bond_name = ["sig2"]+["ls2"]

if (ntriplet>1):
if (ntriplet > 1):
sig2 = random(ntriplet)
ls2 = random(ntriplet)
triplet_mask = np.ones(nspecs**3, dtype=int)
Expand All @@ -81,10 +82,10 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru

sigman = [0.05]

if (nbond>0 and ntriplet>0):
if (nbond > 0 and ntriplet > 0):
hyps = np.hstack([sig1, ls1, sig2, ls2, sigman])
hyps_label = np.hstack([bond_name, triplet_name, ['noise']])
elif (nbond>0):
elif (nbond > 0):
hyps = np.hstack([sig1, ls1, sigman])
hyps_label = np.hstack([bond_name, ['noise']])
else:
Expand All @@ -108,7 +109,7 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru
count = 0
newhyps = []
newlabel = []
if (nbond>1):
if (nbond > 1):
# fix type 0, and only compute type 1 of bonds
hm += [1]
newhyps += [hyps[1]]
Expand All @@ -117,21 +118,21 @@ def generate_hm(nbond, ntriplet, cutoffs=[1, 1], constraint=False, multihyps=Tru
newhyps += [hyps[3]]
newlabel += [hyps_label[3]]
count += 4
elif (nbond>0):
elif (nbond > 0):
# fix sigma, and only vary ls
hm += [1]
newhyps += [hyps[1]]
newlabel += [hyps_label[1]]
count += 2
if (ntriplet>1):
if (ntriplet > 1):
# fix type 0, and only compute type 1 of triplets
hm += [1+count]
newhyps += [hyps[1+count]]
newlabel += [hyps_label[1+count]]
hm += [3+count]
newhyps += [hyps[3+count]]
newlabel += [hyps_label[3+count]]
elif (ntriplet>0):
elif (ntriplet > 0):
# fix sigma, and only vary ls
hm += [1+count]
newhyps += [hyps[1+count]]
Expand Down Expand Up @@ -185,6 +186,7 @@ def get_gp(bodies, kernel_type='mc', multihyps=True) -> GaussianProcess:

return gaussian


def get_params():
parameters = {'unique_species': [2, 1],
'cutoff': 0.8,
Expand All @@ -193,6 +195,7 @@ def get_params():
'db_pts': 30}
return parameters


def get_tstp() -> AtomicEnvironment:
"""Create test point for kernel to compare against"""
# params
Expand All @@ -209,6 +212,67 @@ def get_tstp() -> AtomicEnvironment:
return test_pt


def generate_mb_envs(cutoffs, cell, delt, d1, mask=None):
positions0 = np.array([[0., 0., 0.],
[0., 0.3, 0.],
[0.3, 0., 0.],
[0.0, 0., 0.3],
[1., 1., 0.]])
positions0[1:] += 0.1*np.random.random([4, 3])
triplet = [1, 1, 2, 1]
np.random.shuffle(triplet)
species_1 = np.hstack([triplet, randint(1, 2)])
return generate_mb_envs_pos(positions0, species_1, cutoffs, cell, delt, d1, mask)


def generate_mb_twin_envs(cutoffs, cell, delt, d1, mask=None):

positions0 = np.array([[0., 0., 0.],
[0., 0.3, 0.],
[0.3, 0., 0.],
[0.0, 0., 0.3],
[1., 1., 0.]])
positions0[1:] += 0.1*np.random.random([4, 3])

species_1 = np.array([1, 2, 1, 1, 1], dtype=np.int)
env1 = generate_mb_envs_pos(
positions0, species_1, cutoffs, cell, delt, d1, mask)

positions1 = positions0[[0, 2, 3, 4]]
species_2 = species_1[[0, 2, 3, 4]]
env2 = generate_mb_envs_pos(
positions1, species_2, cutoffs, cell, delt, d1, mask)

return env1, env2


def generate_mb_envs_pos(positions0, species_1, cutoffs, cell, delt, d1, mask=None):

positions = [positions0]

noa = len(positions0)

positions_2 = deepcopy(positions0)
positions_2[0][d1-1] = delt
positions += [positions_2]

positions_3 = deepcopy(positions[0])
positions_3[0][d1-1] = -delt
positions += [positions_3]

test_struc = []
for i in range(3):
test_struc += [struc.Structure(cell, species_1, positions[i])]

env_0 = []
env_p = []
env_m = []
for i in range(noa):
env_0 += [env.AtomicEnvironment(test_struc[0], i, cutoffs)]
env_p += [env.AtomicEnvironment(test_struc[1], i, cutoffs)]
env_m += [env.AtomicEnvironment(test_struc[2], i, cutoffs)]

return [env_0, env_p, env_m]
def generate_envs(cutoffs, delta):
"""
create environment with perturbation on
Expand Down
Loading

0 comments on commit f907c9a

Please sign in to comment.