diff --git a/flare/kernels/mc_simple.py b/flare/kernels/mc_simple.py index 6c87eaf8a..b153c0ab3 100644 --- a/flare/kernels/mc_simple.py +++ b/flare/kernels/mc_simple.py @@ -622,6 +622,277 @@ def two_plus_three_efs_self( return two_e + three_e, two_f + three_f, two_s + three_s +# ----------------------------------------------------------------------------- +# two plus many body kernels +# ----------------------------------------------------------------------------- + + +def two_plus_many_body_mc( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + d2: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+many body kernel. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + d2 (int): Force component of the second environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2b, ls2b, + sigmb, lsmb, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and many-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + sig2 = hyps[0] + ls2 = hyps[1] + sigm = hyps[2] + lsm = hyps[3] + + r_cut_2 = cutoffs[0] + + two_term = two_body_mc_jit( + env1.bond_array_2, + env1.ctype, + env1.etypes, + env2.bond_array_2, + env2.ctype, + env2.etypes, + d1, + d2, + sig2, + ls2, + r_cut_2, + cutoff_func, + ) + + many_term = many_body_mc_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env2.q_neigh_array, + env1.q_neigh_grads, + env2.q_neigh_grads, + env1.ctype, + env2.ctype, + env1.etypes_mb, + env2.etypes_mb, + env1.unique_species, + env2.unique_species, + d1, + d2, + sigm, + lsm, + ) + + return two_term + many_term + + +def two_plus_many_body_mc_grad( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + d2: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+many-body single-element kernel between two force components. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + d2 (int): Force component of the second environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1, + sig2, ls2, sig3, ls3, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + sig2 = hyps[0] + ls2 = hyps[1] + sigm = hyps[2] + lsm = hyps[3] + + r_cut_2 = cutoffs[0] + + kern2, grad2 = two_body_mc_grad_jit( + env1.bond_array_2, + env1.ctype, + env1.etypes, + env2.bond_array_2, + env2.ctype, + env2.etypes, + d1, + d2, + sig2, + ls2, + r_cut_2, + cutoff_func, + ) + + kern_many, gradm = many_body_mc_grad_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env2.q_neigh_array, + env1.q_neigh_grads, + env2.q_neigh_grads, + env1.ctype, + env2.ctype, + env1.etypes_mb, + env2.etypes_mb, + env1.unique_species, + env2.unique_species, + d1, + d2, + sigm, + lsm, + ) + + return kern2 + kern_many, np.hstack([grad2, gradm]) + + +def two_plus_many_body_mc_force_en( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+many-body multi-element kernel between two force and energy + components. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1, + sig2, ls2, sig3, ls3, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + sig2 = hyps[0] + ls2 = hyps[1] + sigm = hyps[2] + lsm = hyps[3] + + r_cut_2 = cutoffs[0] + + two_term = ( + two_body_mc_force_en_jit( + env1.bond_array_2, + env1.ctype, + env1.etypes, + env2.bond_array_2, + env2.ctype, + env2.etypes, + d1, + sig2, + ls2, + r_cut_2, + cutoff_func, + ) + / 2 + ) + + many_term = many_body_mc_force_en_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env1.q_neigh_grads, + env1.ctype, + env2.ctype, + env1.etypes_mb, + env1.unique_species, + env2.unique_species, + d1, + sigm, + lsm, + ) + + return two_term + many_term + + +def two_plus_many_body_mc_en( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+3+many-body single-element energy kernel. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2b ls2b, + sigmb, lsmb, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + sig2 = hyps[0] + ls2 = hyps[1] + sigm = hyps[2] + lsm = hyps[3] + + r_cut_2 = cutoffs[0] + + two_term = ( + two_body_mc_en_jit( + env1.bond_array_2, + env1.ctype, + env1.etypes, + env2.bond_array_2, + env2.ctype, + env2.etypes, + sig2, + ls2, + r_cut_2, + cutoff_func, + ) + / 4 + ) + + many_term = many_body_mc_en_jit( + env1.q_array, + env2.q_array, + env1.ctype, + env2.ctype, + env1.unique_species, + env2.unique_species, + sigm, + lsm, + ) + + return two_term + many_term + + # ----------------------------------------------------------------------------- # two plus three plus many body kernels # ----------------------------------------------------------------------------- @@ -662,7 +933,6 @@ def two_plus_three_plus_many_body_mc( r_cut_2 = cutoffs[0] r_cut_3 = cutoffs[1] - r_cut_m = cutoffs[2] two_term = two_body_mc_jit( env1.bond_array_2, @@ -757,7 +1027,6 @@ def two_plus_three_plus_many_body_mc_grad( r_cut_2 = cutoffs[0] r_cut_3 = cutoffs[1] - r_cut_m = cutoffs[2] kern2, grad2 = two_body_mc_grad_jit( env1.bond_array_2, @@ -4068,7 +4337,6 @@ def many_body_mc_jit( d2 (int): Force component of the second environment. sig (float): many-body signal variance hyperparameter. ls (float): many-body length scale hyperparameter. - r_cut (float): many-body cutoff radius. cutoff_func (Callable): Cutoff function. Return: @@ -4209,7 +4477,6 @@ def many_body_mc_grad_jit( list(set(species1).intersection(set(species2))), dtype=np.int8 ) - for s in useful_species: s1 = np.where(species1 == s)[0][0] s2 = np.where(species2 == s)[0][0] @@ -4452,4 +4719,18 @@ def many_body_mc_en_jit(q_array_1, q_array_2, c1, c2, species1, species2, sig, l "2+3+many_efs_energy": "not implemented", "2+3+many_efs_force": "not implemented", "2+3+many_efs_self": "not implemented", + "two_plus_many_body": two_plus_many_body_mc, + "two_plus_many_body_grad": two_plus_many_body_mc_grad, + "two_plus_many_body_en": two_plus_many_body_mc_en, + "two_plus_many_body_force_en": two_plus_many_body_mc_force_en, + "two_plus_many_body_efs_self": "not implemented", + "two_plus_many_body_efs_force": "not implemented", + "two_plus_many_body_efs_energy": "not implemented", + "2+many": two_plus_many_body_mc, + "2+many_grad": two_plus_many_body_mc_grad, + "2+many_en": two_plus_many_body_mc_en, + "2+many_force_en": two_plus_many_body_mc_force_en, + "2+many_efs_self": "not implemented", + "2+many_efs_force": "not implemented", + "2+many_efs_energy": "not implemented", } diff --git a/flare/kernels/sc.py b/flare/kernels/sc.py index b61cd0d8b..872a4d391 100644 --- a/flare/kernels/sc.py +++ b/flare/kernels/sc.py @@ -22,6 +22,13 @@ * two_plus_three_en, * two_plus_three_force_en +* Two plus many body: + + * two_plus_many_body, + * two_plus_many_body_grad, + * two_plus_many_body_en, + * two_plus_many_body_force_en + * Two plus three plus many body: * two_plus_three_plus_many_body, @@ -283,6 +290,213 @@ def two_plus_three_en(env1, env2, hyps, cutoffs, cutoff_func=cf.quadratic_cutoff return two_term + three_term +# ----------------------------------------------------------------------------- +# two plus many body kernels +# ----------------------------------------------------------------------------- + + +def two_plus_many_body( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + d2: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+many-body single-element kernel between two force components. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + d2 (int): Force component of the second environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2b, ls2b, + sigmb, lsmb, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + two_term = two_body_jit( + env1.bond_array_2, + env2.bond_array_2, + d1, + d2, + hyps[0], + hyps[1], + cutoffs[0], + cutoff_func, + ) + + many_term = many_body_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env2.q_neigh_array, + env1.q_neigh_grads, + env2.q_neigh_grads, + d1, + d2, + hyps[2], + hyps[3], + ) + + return two_term + many_term + + +def two_plus_many_body_grad( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + d2: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+many-body single-element kernel between two force components. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + d2 (int): Force component of the second environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2b, ls2b, + sigmb, lsmb, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + kern2, ls2, sig2 = two_body_grad_jit( + env1.bond_array_2, + env2.bond_array_2, + d1, + d2, + hyps[0], + hyps[1], + cutoffs[0], + cutoff_func, + ) + + kern_many, sigm, lsm = many_body_grad_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env2.q_neigh_array, + env1.q_neigh_grads, + env2.q_neigh_grads, + d1, + d2, + hyps[2], + hyps[3], + ) + + return kern2 + kern_many, np.array([sig2, ls2, sigm, lsm]) + + +def two_plus_many_body_force_en( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + d1: int, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+3+many-body single-element kernel between two force and energy components. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + d1 (int): Force component of the first environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2, ls2, + sig3, ls3, sigm, lsm, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + two_term = ( + two_body_force_en_jit( + env1.bond_array_2, + env2.bond_array_2, + d1, + hyps[0], + hyps[1], + cutoffs[0], + cutoff_func, + ) + / 2 + ) + + many_term = many_body_force_en_jit( + env1.q_array, + env2.q_array, + env1.q_neigh_array, + env1.q_neigh_grads, + d1, + hyps[2], + hyps[3], + ) + + return two_term + many_term + + +def two_plus_many_body_en( + env1: AtomicEnvironment, + env2: AtomicEnvironment, + hyps, + cutoffs, + cutoff_func=cf.quadratic_cutoff, +): + """2+3+many-body single-element energy kernel. + + Args: + env1 (AtomicEnvironment): First local environment. + env2 (AtomicEnvironment): Second local environment. + hyps (np.ndarray): Hyperparameters of the kernel function (sig2, ls2, + sig3, ls3, sigm, lsm, sig_n). + cutoffs (np.ndarray): Two-element array containing the 2- and 3-body + cutoffs. + cutoff_func (Callable): Cutoff function of the kernel. + + Return: + float: Value of the 2+3+many-body kernel. + """ + + two_term = two_body_en_jit( + env1.bond_array_2, env2.bond_array_2, hyps[0], hyps[1], cutoffs[0], cutoff_func + ) + + three_term = three_body_en_jit( + env1.bond_array_3, + env2.bond_array_3, + env1.cross_bond_inds, + env2.cross_bond_inds, + env1.cross_bond_dists, + env2.cross_bond_dists, + env1.triplet_counts, + env2.triplet_counts, + hyps[2], + hyps[3], + cutoffs[1], + cutoff_func, + ) + + many_term = many_body_en_jit(env1.q_array, env2.q_array, hyps[4], hyps[5]) + + return two_term + three_term + many_term + + # ----------------------------------------------------------------------------- # two plus three plus many body kernels # ----------------------------------------------------------------------------- @@ -2116,6 +2330,13 @@ def triplet_force_en_kernel( "many_efs_energy": "not implemented", "many_efs_force": "not implemented", "many_efs_self": "not implemented", + "two_plus_many_body": two_plus_many_body, + "two_plus_many_body_grad": two_plus_many_body_grad, + "two_plus_many_body_en": two_plus_many_body_en, + "two_plus_many_body_force_en": two_plus_many_body_force_en, + "two_plus_many_body_efs_self": "not implemented", + "two_plus_many_body_efs_force": "not implemented", + "two_plus_many_body_efs_energy": "not implemented", "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, @@ -2127,6 +2348,13 @@ def triplet_force_en_kernel( "2+3+many_efs_energy": "not implemented", "2+3+many_efs_force": "not implemented", "2+3+many_efs_self": "not implemented", + "2+many": two_plus_many_body, + "2+many_grad": two_plus_many_body_grad, + "2+many_en": two_plus_many_body_en, + "2+many_force_en": two_plus_many_body_force_en, + "2+many_efs_self": "not implemented", + "2+many_efs_force": "not implemented", + "2+many_efs_energy": "not implemented", } diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 8a818e562..8fe5a2b07 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -10,7 +10,7 @@ from .fake_gp import generate_mb_envs -list_to_test = [["2"], ["3"], ["2", "3"], ["2", "3", "many"]] +list_to_test = [["2"], ["3"], ["2", "3"], ["2", "3", "many"],['2','many']] list_type = ["sc", "mc"]