<commit_before>"""
Draft of a force-field. At this point only a Lennard Jones term is implemented.
"""
from scipy.spatial import cKDTree
from numba import jit
from ..utils.geometry import dist
from ..utils.constants import par_s_ij, par_eps_ij


def compute_neighbors(coords, exclusions, cut_off):
    """
    Use a KD-tree (from scipy) to compute the neighbors atoms for each atom.

    Parameters
    ----------
    coords : array (m, 3)
        Cartesian coordinates of a molecule
    exclusions : set of tuples
        Pairs of atoms excluded from the computation of the neighbors
    cut_off : float
        Only pairs of atoms closer than cut_off will be used to compute the
        neighbors.

    Results
    -------
    neighbors: set of tuples
        Pairs of neighbors atoms within a given "cut_off" and excluding
        "exclusions"
    """
    tree_c = cKDTree(coords)
    all_pairs = tree_c.query_pairs(cut_off)
    return all_pairs - exclusions


def LJ(neighbors, xyz, elements, mode='A'):
    """
    Lennard Jones energy term

    .. math::

    LJ_{ij} = \\epsilon \left [ \left (\\frac{\sigma_{ij}}{r_{ij}} \\right)^{12}
     - 2 \left (\\frac{\sigma_{ij}}{r_{ij}} \\right)^{6} \\right]

    \sigma_{ij} is the distance at which the potential reaches its minimum
    \\epsilon_{ij} is the depth of the potential well
    r_{ij} is the distance between the particles

    Parameters
    ----------
    XXX : XXX
        XXX
    XXX : XXX
        XXX<commit_msg>more modification of LJ equation<commit_after>"""
Draft of a force-field. At this point only a Lennard Jones term is implemented.
"""
from scipy.spatial import cKDTree
from numba import jit
from ..utils.geometry import dist
from ..utils.constants import par_s_ij, par_eps_ij


def compute_neighbors(coords, exclusions, cut_off):
    """
    Use a KD-tree (from scipy) to compute the neighbors atoms for each atom.

    Parameters
    ----------
    coords : array (m, 3)
        Cartesian coordinates of a molecule
    exclusions : set of tuples
        Pairs of atoms excluded from the computation of the neighbors
    cut_off : float
        Only pairs of atoms closer than cut_off will be used to compute the
        neighbors.

    Results
    -------
    neighbors: set of tuples
        Pairs of neighbors atoms within a given "cut_off" and excluding
        "exclusions"
    """
    tree_c = cKDTree(coords)
    all_pairs = tree_c.query_pairs(cut_off)
    return all_pairs - exclusions


def LJ(neighbors, xyz, elements, mode='A'):
    """
    Lennard Jones energy term

    .. math::

    LJ_{ij} = \\epsilon \\left [ \\left (\\frac{\\sigma_{ij}}{r_{ij}} \\right)^{12}
     - 2 \\left (\\frac{\\sigma_{ij}}{r_{ij}} \\right)^{6} \\right]

    \\sigma_{ij} is the distance at which the potential reaches its minimum
    \\epsilon_{ij} is the depth of the potential well
    r_{ij} is the distance between the particles

    Parameters
    ----------
    XXX : XXX
        XXX
    XXX : XXX
        XXX
{'fraction_non_alphanumeric': 0.10772833723653395, 'fraction_numerical': 0.0033456005352960855, 'mean_word_length': 2.4486735870818914, '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': '7207432', 'n_tokens_mistral': 926, 'n_tokens_neox': 891, 'n_words': 335}