<commit_before># distutils: language = c++
#cython: embedsignature=True

from libcpp.string cimport string
from libcpp.vector cimport vector

from functools import total_ordering
import itertools

from tools.geometry import distance


cdef extern from "dataClasses.hpp" namespace "bude":
    cdef cppclass AtomDataFF:
        AtomDataFF(string, string, double, double, double, double, double, double, double) except +
        string atomName, electrostaticType;
        double radius, hydrophobicPotential, hardness, distNpNp, distNpP, radiusScale, electrostaticPotential


cdef extern from "calculateEnergies.cpp" namespace "bude":
    cdef vector[double] calculatePairEnergy(AtomDataFF *, AtomDataFF *, double)


cdef class PyAtomData:
    cdef AtomDataFF *thisptr
    def __cinit__(self, string atomName, string electrostaticType, double radius, double hydrophobicPotential,
                  double hardness, double distNpNp, double distNpP, double radiusScale, double electrostaticPotential):
        self.thisptr = new AtomDataFF(atomName, electrostaticType, radius, hydrophobicPotential, hardness, distNpNp,
                                      distNpP, radiusScale, electrostaticPotential)
    def __dealloc__(self):
        del self.thisptr


cpdef find_buff_interactions(ampal, ff, internal=False):
    """Finds BUFF interactions for a given ampal object and force field.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object with a get_atoms method.
    ff: BuffForceField
        The force field is used to derive the minimum interaction distance.

    Returns
    -------
    interaction_set: [(Atom, Atom)]
        All of the atom pairs in range of interacting in BUFF but not within
        covalent bond distance.
    """
    cdef double ffco, m_dist
    ffco = ff.distance_cutoff
    cen_mons = []
    for monomer in ampal.get_monomers():
        if hasattr(monomer, 'reference_atom'):
            cen_mons.append((monomer, monomer.atoms[monomer.reference_atom]))
        # TODO: Add centre of mass as an option here. Needs careful thought so not to miss interactions.
    mon_pairs = []
    interactions = []
    ncaco = ['N', 'CA', 'C', 'O']
    for (monomer_a, ref_atom_a), (monomer_b, ref_atom_b) in itertools.combinations(cen_mons, 2):
        #### Could do this with filters
        if not internal:
            if monomer_a.ampal_parent != monomer_b.ampal_parent:
                m_dist = distance(ref_atom_a, ref_atom_b)
            else:
                continue
        ####
        else:
            m_dist = distance(ref_atom_a, ref_atom_b)
        if m_dist <= ffco:
            a_atoms = [atom for atom in monomer_a.atoms.values() if atom._ff_id is not None]
            b_atoms = [atom for atom in monomer_b.atoms.values() if atom._ff_id is not None]
            if not internal:
                interactions.extend(itertools.product(a_atoms, b_atoms))
            else:
                for interaction in itertools.product(a_atoms, b_atoms):
                    chain1, res1 = interaction[0].unique_id[:2]
                    chain2, res2 = interaction[1].unique_id[:2]
                    if (interaction[0].res_label in ncaco) and (interaction[1].res_label in ncaco):
                        if (chain2, int(res2)) == (chain1, int(res1) + 1):
                            pass
                        else:
                            interactions.append(interaction)
                    else:
                        interactions.append(interaction)
    return interactions


cpdef get_within_ff_cutoff(interaction_pairs, double force_field_cutoff):
    """Finds BUFF interactions for a given ampal object and force field.

    Parameters
    ----------
    interaction_pairs: [Monomer]
        Any AMPAL object with a get_atoms method.
    force_field_cutoff: float
        Minimum interaction distance.

    Returns
    -------
    interaction_set: [(Atom, Atom)]
        All of the atom pairs in range of interacting in BUFF but not within
        covalent bond distance.
    """
    cdef double m_dist
    interactions = []
    for ref_atom_a, ref_atom_b in interaction_pairs:
        m_dist = distance(ref_atom_a._vector, ref_atom_b._vector)
        if m_dist <= force_field_cutoff:
            a_atoms = [atom for atom in ref_atom_a.ampal_parent.atoms.values() if atom._ff_id is not None]
            b_atoms = [atom for atom in ref_atom_b.ampal_parent.atoms.values() if atom._ff_id is not None]
            interactions.extend(itertools.product(a_atoms, b_atoms))
    return interactions


cpdef score_interactions(interactions, ff, threshold = 1.1):
    """Scores a set of interactions.

    Parameters
    ----------
    interactions: [(Atom, Atom)]
        An interaction is defined as a pair of Atoms that have had force field
        parameters assigned to them. This means that they should have a
        'ff_params' tag in their tags.
    ff: BuffForceField
        The force field used for scoring.
    threshold: float
        Cutoff distance for assigning interactions that are covalent bonds.

    Returns
    -------
    buff_score: BuffScore
        A BuffScore object with information about each of the interactions and
        the atoms involved.
    """
    cdef PyAtomData a_params, b_params
    cdef double dist, ffco
    scores = []
    ffco = ff.distance_cutoff
    ffpsd = ff.parameter_struct_dict
    for a, b in interactions:
        a_params = ffpsd[a._ff_id[0]][a._ff_id[1]]
        b_params = ffpsd[b._ff_id[0]][b._ff_id[1]]
        dist = distance(a._vector, b._vector)
        if dist <= ffco:
            scores.append(calculatePairEnergy(a_params.thisptr, b_params.thisptr, dist))
    buff_score = BuffScore(interactions, scores)
    return buff_score


def score_ampal(ampal, ff, threshold=1.1, internal=False):
    """Finds interactions in an AMPAL object and returns the BUFF score.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object that inherits from BaseAmpal.
    ff: BuffForceField
        Force field for BUFF.
    threshold: float
        Cutoff distance for assigning interactions that are covalent bonds.
    internal: bool
        Measures internal energy if true.

    Returns
    -------
    BUFF_score: BUFFScore
        A BUFFScore object with information about each of the interactions and
        the atoms involved.

    """
    interactions = find_buff_interactions(ampal, ff, internal=internal)
    return score_interactions(interactions, ff, threshold=threshold)


def score_intra_ampal(ampal, ff, no_neighbour_backbone=True, backbone_atoms=('N', 'CA', 'C', 'O')):
    """Finds interactions within an AMPAL object and returns the BUFF score.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object that inherits from BaseAmpal.
    ff: BuffForceField
        Force field for BUFF.
    no_neighbour_backbone: bool
        Ignore interactions between neighbouring backbone atoms.
    backbone_atoms: (str)
        A tuple containing labels for the backbone atoms in the polymer/s.

    Returns
    -------
    BUFF_score: BUFFScore
        A BUFFScore object with information about each of the interactions and
        the atoms involved.

    """
    grp = [mon[mon.reference_atom] for mon in ampal.get_monomers() if hasattr(mon, 'reference_atom')]
    gross_interactions = itertools.combinations(grp, 2)
    interactions = get_within_ff_cutoff(gross_interactions, ff.distance_cutoff)
    if no_neighbour_backbone:
        interactions = filter(lambda x: check_if_backbone_neighbours(x, backbone_atoms), interactions)
    return score_interactions(interactions, ff)
<commit_msg>Removed references to `threshold`, as this is no longer used.<commit_after># distutils: language = c++
#cython: embedsignature=True

from libcpp.string cimport string
from libcpp.vector cimport vector

from functools import total_ordering
import itertools

from tools.geometry import distance


cdef extern from "dataClasses.hpp" namespace "bude":
    cdef cppclass AtomDataFF:
        AtomDataFF(string, string, double, double, double, double, double, double, double) except +
        string atomName, electrostaticType;
        double radius, hydrophobicPotential, hardness, distNpNp, distNpP, radiusScale, electrostaticPotential


cdef extern from "calculateEnergies.cpp" namespace "bude":
    cdef vector[double] calculatePairEnergy(AtomDataFF *, AtomDataFF *, double)


cdef class PyAtomData:
    cdef AtomDataFF *thisptr
    def __cinit__(self, string atomName, string electrostaticType, double radius, double hydrophobicPotential,
                  double hardness, double distNpNp, double distNpP, double radiusScale, double electrostaticPotential):
        self.thisptr = new AtomDataFF(atomName, electrostaticType, radius, hydrophobicPotential, hardness, distNpNp,
                                      distNpP, radiusScale, electrostaticPotential)
    def __dealloc__(self):
        del self.thisptr


cpdef find_buff_interactions(ampal, ff, internal=False):
    """Finds BUFF interactions for a given ampal object and force field.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object with a get_atoms method.
    ff: BuffForceField
        The force field is used to derive the minimum interaction distance.

    Returns
    -------
    interaction_set: [(Atom, Atom)]
        All of the atom pairs in range of interacting in BUFF but not within
        covalent bond distance.
    """
    cdef double ffco, m_dist
    ffco = ff.distance_cutoff
    cen_mons = []
    for monomer in ampal.get_monomers():
        if hasattr(monomer, 'reference_atom'):
            cen_mons.append((monomer, monomer.atoms[monomer.reference_atom]))
        # TODO: Add centre of mass as an option here. Needs careful thought so not to miss interactions.
    mon_pairs = []
    interactions = []
    ncaco = ['N', 'CA', 'C', 'O']
    for (monomer_a, ref_atom_a), (monomer_b, ref_atom_b) in itertools.combinations(cen_mons, 2):
        #### Could do this with filters
        if not internal:
            if monomer_a.ampal_parent != monomer_b.ampal_parent:
                m_dist = distance(ref_atom_a, ref_atom_b)
            else:
                continue
        ####
        else:
            m_dist = distance(ref_atom_a, ref_atom_b)
        if m_dist <= ffco:
            a_atoms = [atom for atom in monomer_a.atoms.values() if atom._ff_id is not None]
            b_atoms = [atom for atom in monomer_b.atoms.values() if atom._ff_id is not None]
            if not internal:
                interactions.extend(itertools.product(a_atoms, b_atoms))
            else:
                for interaction in itertools.product(a_atoms, b_atoms):
                    chain1, res1 = interaction[0].unique_id[:2]
                    chain2, res2 = interaction[1].unique_id[:2]
                    if (interaction[0].res_label in ncaco) and (interaction[1].res_label in ncaco):
                        if (chain2, int(res2)) == (chain1, int(res1) + 1):
                            pass
                        else:
                            interactions.append(interaction)
                    else:
                        interactions.append(interaction)
    return interactions


cpdef get_within_ff_cutoff(interaction_pairs, double force_field_cutoff):
    """Finds BUFF interactions for a given ampal object and force field.

    Parameters
    ----------
    interaction_pairs: [Monomer]
        Any AMPAL object with a get_atoms method.
    force_field_cutoff: float
        Minimum interaction distance.

    Returns
    -------
    interaction_set: [(Atom, Atom)]
        All of the atom pairs in range of interacting in BUFF but not within
        covalent bond distance.
    """
    cdef double m_dist
    interactions = []
    for ref_atom_a, ref_atom_b in interaction_pairs:
        m_dist = distance(ref_atom_a._vector, ref_atom_b._vector)
        if m_dist <= force_field_cutoff:
            a_atoms = [atom for atom in ref_atom_a.ampal_parent.atoms.values() if atom._ff_id is not None]
            b_atoms = [atom for atom in ref_atom_b.ampal_parent.atoms.values() if atom._ff_id is not None]
            interactions.extend(itertools.product(a_atoms, b_atoms))
    return interactions


cpdef score_interactions(interactions, ff):
    """Scores a set of interactions.

    Parameters
    ----------
    interactions: [(Atom, Atom)]
        An interaction is defined as a pair of Atoms that have had force field
        parameters assigned to them. This means that they should have a
        'ff_params' tag in their tags.
    ff: BuffForceField
        The force field used for scoring.

    Returns
    -------
    buff_score: BuffScore
        A BuffScore object with information about each of the interactions and
        the atoms involved.
    """
    cdef PyAtomData a_params, b_params
    cdef double dist, ffco
    scores = []
    ffco = ff.distance_cutoff
    ffpsd = ff.parameter_struct_dict
    for a, b in interactions:
        a_params = ffpsd[a._ff_id[0]][a._ff_id[1]]
        b_params = ffpsd[b._ff_id[0]][b._ff_id[1]]
        dist = distance(a._vector, b._vector)
        if dist <= ffco:
            scores.append(calculatePairEnergy(a_params.thisptr, b_params.thisptr, dist))
    buff_score = BuffScore(interactions, scores)
    return buff_score


def score_ampal(ampal, ff, internal=False):
    """Finds interactions in an AMPAL object and returns the BUFF score.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object that inherits from BaseAmpal.
    ff: BuffForceField
        Force field for BUFF.
    internal: bool
        Measures internal energy if true.

    Returns
    -------
    BUFF_score: BUFFScore
        A BUFFScore object with information about each of the interactions and
        the atoms involved.

    """
    interactions = find_buff_interactions(ampal, ff, internal=internal)
    return score_interactions(interactions, ff)


def score_intra_ampal(ampal, ff, no_neighbour_backbone=True, backbone_atoms=('N', 'CA', 'C', 'O')):
    """Finds interactions within an AMPAL object and returns the BUFF score.

    Parameters
    ----------
    ampal: AMPAL object
        Any AMPAL object that inherits from BaseAmpal.
    ff: BuffForceField
        Force field for BUFF.
    no_neighbour_backbone: bool
        Ignore interactions between neighbouring backbone atoms.
    backbone_atoms: (str)
        A tuple containing labels for the backbone atoms in the polymer/s.

    Returns
    -------
    BUFF_score: BUFFScore
        A BUFFScore object with information about each of the interactions and
        the atoms involved.

    """
    grp = [mon[mon.reference_atom] for mon in ampal.get_monomers() if hasattr(mon, 'reference_atom')]
    gross_interactions = itertools.combinations(grp, 2)
    interactions = get_within_ff_cutoff(gross_interactions, ff.distance_cutoff)
    if no_neighbour_backbone:
        interactions = filter(lambda x: check_if_backbone_neighbours(x, backbone_atoms), interactions)
    return score_interactions(interactions, ff)

{'fraction_non_alphanumeric': 0.08195961602118504, 'fraction_numerical': 0.0030453492221118833, 'mean_word_length': 2.6018121125417264, 'pattern_counts': {'":': 4, '<': 9, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6423948', 'n_tokens_mistral': 4560, 'n_tokens_neox': 4340, 'n_words': 1345}