<commit_before>import numpy as np
from numba import njit
from math import exp, floor
from typing import Callable

from flare.kernels.cutoffs import quadratic_cutoff

from time import time

def grid_kernel_sephyps(kern_type,
                  data, grids, fj, fdj,
                  c2, etypes2, 
                  cutoff_2b, cutoff_3b, cutoff_mb,
                  nspec, spec_mask,
                  nbond, bond_mask,
                  ntriplet, triplet_mask,
                  ncut3b, cut3b_mask,
                  nmb, mb_mask,
                  sig2, ls2, sig3, ls3, sigm, lsm,
                  cutoff_func=quadratic_cutoff):
    '''
    Args:
        data: a single env of a list of envs
    '''

    bc1 = spec_mask[c2]
    bc2 = spec_mask[etypes2[0]]
    bc3 = spec_mask[etypes2[1]]
    ttype = triplet_mask[nspec * nspec * bc1 + nspec*bc2 + bc3]
    ls = ls3[ttype]
    sig = sig3[ttype]
    cutoffs = [cutoff_2b, cutoff_3b]

    hyps = [sig, ls]
    return grid_kernel(kern_type,
                   data, grids, fj, fdj,
                   c2, etypes2, 
                   hyps, cutoffs, cutoff_func)


def grid_kernel(kern_type,
                struc, grids, fj, fdj,
                c2, etypes2, 
                hyps: 'ndarray', cutoffs,
                cutoff_func: Callable = quadratic_cutoff):

    r_cut = cutoffs[1]

    if not isinstance(struc, list):
        struc = [struc]

    kern = 0
    for env in struc:
        kern += grid_kernel_env(kern_type,
                    env, grids, fj, fdj,
                    c2, etypes2, 
                    hyps, r_cut, cutoff_func)

    return kern


def grid_kernel_env(kern_type,
                env1, grids, fj, fdj,
                c2, etypes2, 
                hyps: 'ndarray', r_cut: float,
                cutoff_func: Callable = quadratic_cutoff):

    # pre-compute constants that appear in the inner loop
    sig = hyps[0]
    ls = hyps[1]
    derivative = derv_dict[kern_type]

    # collect all the triplets in this training env
    triplet_coord_list = get_triplets_for_kern(env1.bond_array_3, env1.ctype, env1.etypes,
        env1.cross_bond_inds, env1.cross_bond_dists, env1.triplet_counts,
        c2, etypes2)

    if len(triplet_coord_list) == 0: # no triplets
        if derivative:
            return np.zeros((3, grids.shape[0]), dtype=np.float64)
        else:
            return np.zeros(grids.shape[0], dtype=np.float64)

    triplet_coord_list = np.array(triplet_coord_list)
    triplet_list = triplet_coord_list[:, :3] # (n_triplets, 3)
    coord_list = triplet_coord_list[:, 3:] # ((n_triplets, 9)

    # calculate distance difference & exponential part
    ls1 = 1 / (2 * ls * ls)
    D = 0
    rij_list = []
    for r in range(3):
        rj, ri = np.meshgrid(grids[:, r], triplet_list[:, r])
        rij = ri - rj
        D += rij * rij # (n_triplets, n_grids)
        rij_list.append(rij)

    kern_exp = (sig * sig) * np.exp(- D * ls1)

    # calculate cutoff of the triplets
    fi, fdi = triplet_cutoff(triplet_list, r_cut, coord_list, derivative,
        cutoff_func) # (n_triplets, 1)

    # calculate the derivative part
    kern_func = kern_dict[kern_type]
    kern = kern_func(kern_exp, fi, fj, fdi, fdj,
             rij_list, coord_list, ls)

    return kern


def en_en(kern_exp, fi, fj, *args):
    '''energy map + energy block'''
    fifj = fi @ fj.T # (n_triplets, n_grids)
    kern = np.sum(kern_exp * fifj, axis=0) / 9 # (n_grids,)
    return kern


def en_force(kern_exp, fi, fj, fdi, fdj, 
             rij_list, coord_list, ls):
    '''energy map + force block'''
    fifj = fi @ fj.T # (n_triplets, n_grids)
    ls2 = 1 / (ls * ls)
    n_trplt, n_grids = kern_exp.shape
    kern = np.zeros((3, n_grids), dtype=np.float64)
    for d in range(3):
        B = 0
        fdij = fdi[:, [d]] @ fj.T
        for r in range(3):
            rij = rij_list[r]
            # column-wise multiplication
            # coord_list[:, [r]].shape = (n_triplets, 1)
            B += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids)

        kern[d, :] = - np.sum(kern_exp * (B * ls2 * fifj + fdij), axis=0) / 3 # (n_grids,)
    return kern


def force_en(kern_exp, fi, fj, fdi, fdj, 
             rij_list, coord_list, ls):
    '''force map + energy block'''
    ls2 = 1 / (ls * ls)
    fifj = fi @ fj.T # (n_triplets, n_grids)
    fdji = fi @ fdj.T
    # only r = 0 is non zero, since the grid coords are all (1, 0, 0)
    B = rij_list[0] # (n_triplets, n_grids)
    kern = np.sum(kern_exp * (B * ls2 * fifj - fdji), axis=0) / 3 # (n_grids,)
    return kern


def force_force(kern_exp, fi, fj, fdi, fdj, 
                rij_list, coord_list, ls):
    '''force map + force block'''
    ls2 = 1 / (ls * ls)
    ls3 = ls2 * ls2 
    
    n_trplt, n_grids = kern_exp.shape
    kern = np.zeros((3, n_grids), dtype=np.float64)

    fifj = fi @ fj.T # (n_triplets, n_grids)
    fdji = (fi * ls2) @ fdj.T

    B = rij_list[0] * ls3 # B and C both have opposite signs with that in the three_body_helper_1
    
    for d in range(3):
        fdij = (fdi[:, [d]] * ls2) @ fj.T
        I = fdi[:, [d]] @ fdj.T
        J = rij_list[0] * fdij # (n_triplets, n_grids)

        A = np.repeat(ls2 * coord_list[:, [3*d]], n_grids, axis=1)
        C = 0
        for r in range(3):
            rij = rij_list[r]
            # column-wise multiplication
            # coord_list[:, [r]].shape = (n_triplets, 1)
            C += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids)

        IJKL = I - J + C * fdji + (A - B * C) * fifj
        kern[d, :] = np.sum(kern_exp * IJKL, axis=0)

    return kern



def triplet_cutoff(triplets, r_cut, coords, derivative=False, cutoff_func=quadratic_cutoff):

    dfj_list = np.zeros((len(triplets), 3), dtype=np.float64)

    if derivative:
        for d in range(3):
            s = 3 * d
            e = 3 * (d + 1)
            f0, df0 = cutoff_func(r_cut, triplets, coords[:, s:e])
            dfj = df0[:, 0] *  f0[:, 1] *  f0[:, 2] + \
                   f0[:, 0] * df0[:, 1] *  f0[:, 2] + \
                   f0[:, 0] *  f0[:, 1] * df0[:, 2]
#            dfj = np.expand_dims(dfj, axis=1)
            dfj_list[:, d] = dfj
    else:
        f0, _ = cutoff_func(r_cut, triplets, 0) # (n_grid, 3)

    fj = f0[:, 0] * f0[:, 1] * f0[:, 2] # (n_grid,)
    fj = np.expand_dims(fj, axis=1)

    return fj, dfj_list


@njit
def get_triplets_for_kern(bond_array_1, c1, etypes1,
                          cross_bond_inds_1, cross_bond_dists_1,
                          triplets_1,
                          c2, etypes2):

    #triplet_list = np.empty((0, 6), dtype=np.float64)
    triplet_list = []<commit_msg>add 'del' to grid_kern<commit_after>import numpy as np
from numba import njit
from math import exp, floor
from typing import Callable

from flare.kernels.cutoffs import quadratic_cutoff

from time import time

def grid_kernel_sephyps(kern_type,
                  data, grids, fj, fdj,
                  c2, etypes2, 
                  cutoff_2b, cutoff_3b, cutoff_mb,
                  nspec, spec_mask,
                  nbond, bond_mask,
                  ntriplet, triplet_mask,
                  ncut3b, cut3b_mask,
                  nmb, mb_mask,
                  sig2, ls2, sig3, ls3, sigm, lsm,
                  cutoff_func=quadratic_cutoff):
    '''
    Args:
        data: a single env of a list of envs
    '''

    bc1 = spec_mask[c2]
    bc2 = spec_mask[etypes2[0]]
    bc3 = spec_mask[etypes2[1]]
    ttype = triplet_mask[nspec * nspec * bc1 + nspec*bc2 + bc3]
    ls = ls3[ttype]
    sig = sig3[ttype]
    cutoffs = [cutoff_2b, cutoff_3b]

    hyps = [sig, ls]
    return grid_kernel(kern_type,
                   data, grids, fj, fdj,
                   c2, etypes2, 
                   hyps, cutoffs, cutoff_func)


def grid_kernel(kern_type,
                struc, grids, fj, fdj,
                c2, etypes2, 
                hyps: 'ndarray', cutoffs,
                cutoff_func: Callable = quadratic_cutoff):

    r_cut = cutoffs[1]

    if not isinstance(struc, list):
        struc = [struc]

    kern = 0
    for env in struc:
        kern += grid_kernel_env(kern_type,
                    env, grids, fj, fdj,
                    c2, etypes2, 
                    hyps, r_cut, cutoff_func)

    return kern


def grid_kernel_env(kern_type,
                env1, grids, fj, fdj,
                c2, etypes2, 
                hyps: 'ndarray', r_cut: float,
                cutoff_func: Callable = quadratic_cutoff):

    # pre-compute constants that appear in the inner loop
    sig = hyps[0]
    ls = hyps[1]
    derivative = derv_dict[kern_type]

    # collect all the triplets in this training env
    triplet_coord_list = get_triplets_for_kern(env1.bond_array_3, env1.ctype, env1.etypes,
        env1.cross_bond_inds, env1.cross_bond_dists, env1.triplet_counts,
        c2, etypes2)

    if len(triplet_coord_list) == 0: # no triplets
        if derivative:
            return np.zeros((3, grids.shape[0]), dtype=np.float64)
        else:
            return np.zeros(grids.shape[0], dtype=np.float64)

    triplet_coord_list = np.array(triplet_coord_list)
    triplet_list = triplet_coord_list[:, :3] # (n_triplets, 3)
    coord_list = triplet_coord_list[:, 3:] # ((n_triplets, 9)
    del triplet_coord_list

    # calculate distance difference & exponential part
    ls1 = 1 / (2 * ls * ls)
    D = 0
    rij_list = []
    for r in range(3):
        rj, ri = np.meshgrid(grids[:, r], triplet_list[:, r])
        rij = ri - rj
        D += rij * rij # (n_triplets, n_grids)
        rij_list.append(rij)

    kern_exp = (sig * sig) * np.exp(- D * ls1)
    del D

    # calculate cutoff of the triplets
    fi, fdi = triplet_cutoff(triplet_list, r_cut, coord_list, derivative,
        cutoff_func) # (n_triplets, 1)
    del triplet_list

    # calculate the derivative part
    kern_func = kern_dict[kern_type]
    kern = kern_func(kern_exp, fi, fj, fdi, fdj,
             rij_list, coord_list, ls)

    return kern


def en_en(kern_exp, fi, fj, *args):
    '''energy map + energy block'''
    fifj = fi @ fj.T # (n_triplets, n_grids)
    kern = np.sum(kern_exp * fifj, axis=0) / 9 # (n_grids,)
    return kern


def en_force(kern_exp, fi, fj, fdi, fdj, 
             rij_list, coord_list, ls):
    '''energy map + force block'''
    fifj = fi @ fj.T # (n_triplets, n_grids)
    ls2 = 1 / (ls * ls)
    n_trplt, n_grids = kern_exp.shape
    kern = np.zeros((3, n_grids), dtype=np.float64)
    for d in range(3):
        B = 0
        fdij = fdi[:, [d]] @ fj.T
        for r in range(3):
            rij = rij_list[r]
            # column-wise multiplication
            # coord_list[:, [r]].shape = (n_triplets, 1)
            B += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids)

        kern[d, :] = - np.sum(kern_exp * (B * ls2 * fifj + fdij), axis=0) / 3 # (n_grids,)
    return kern


def force_en(kern_exp, fi, fj, fdi, fdj, 
             rij_list, coord_list, ls):
    '''force map + energy block'''
    ls2 = 1 / (ls * ls)
    fifj = fi @ fj.T # (n_triplets, n_grids)
    fdji = fi @ fdj.T
    # only r = 0 is non zero, since the grid coords are all (1, 0, 0)
    B = rij_list[0] # (n_triplets, n_grids)
    kern = np.sum(kern_exp * (B * ls2 * fifj - fdji), axis=0) / 3 # (n_grids,)
    return kern


def force_force(kern_exp, fi, fj, fdi, fdj, 
                rij_list, coord_list, ls):
    '''force map + force block'''
    ls2 = 1 / (ls * ls)
    ls3 = ls2 * ls2 
    
    n_trplt, n_grids = kern_exp.shape
    kern = np.zeros((3, n_grids), dtype=np.float64)

    fifj = fi @ fj.T # (n_triplets, n_grids)
    fdji = (fi * ls2) @ fdj.T

    B = rij_list[0] * ls3 # B and C both have opposite signs with that in the three_body_helper_1
    
    for d in range(3):
        fdij = (fdi[:, [d]] * ls2) @ fj.T
        I = fdi[:, [d]] @ fdj.T
        J = rij_list[0] * fdij # (n_triplets, n_grids)

        A = np.repeat(ls2 * coord_list[:, [3*d]], n_grids, axis=1)
        C = 0
        for r in range(3):
            rij = rij_list[r]
            # column-wise multiplication
            # coord_list[:, [r]].shape = (n_triplets, 1)
            C += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids)

        IJKL = I - J + C * fdji + (A - B * C) * fifj
        kern[d, :] = np.sum(kern_exp * IJKL, axis=0)

    return kern



def triplet_cutoff(triplets, r_cut, coords, derivative=False, cutoff_func=quadratic_cutoff):

    dfj_list = np.zeros((len(triplets), 3), dtype=np.float64)

    if derivative:
        for d in range(3):
            s = 3 * d
            e = 3 * (d + 1)
            f0, df0 = cutoff_func(r_cut, triplets, coords[:, s:e])
            dfj = df0[:, 0] *  f0[:, 1] *  f0[:, 2] + \
                   f0[:, 0] * df0[:, 1] *  f0[:, 2] + \
                   f0[:, 0] *  f0[:, 1] * df0[:, 2]
            dfj_list[:, d] = dfj
    else:
        f0, _ = cutoff_func(r_cut, triplets, 0) # (n_grid, 3)

    fj = f0[:, 0] * f0[:, 1] * f0[:, 2] # (n_grid,)
    fj = np.expand_dims(fj, axis=1)

    return fj, dfj_list


@njit
def get_triplets_for_kern(bond_array_1, c1, etypes1,
                          cross_bond_inds_1, cross_bond_dists_1,
                          triplets_1,
                          c2, etypes2):

    #triplet_list = np.empty((0, 6), dtype=np.float64)
    triplet_list = []
{'fraction_non_alphanumeric': 0.12624209575429088, 'fraction_numerical': 0.024616079494128276, 'mean_word_length': 1.9528784174260947, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2201336', 'n_tokens_mistral': 5415, 'n_tokens_neox': 5003, 'n_words': 1254}