<commit_before>from pyemma.util.annotators import deprecated
__author__ = 'noe'

import mdtraj
from mdtraj.geometry.dihedral import _get_indices_phi, \
    _get_indices_psi, compute_dihedrals

import numpy as np
import warnings

from pyemma.util.log import getLogger

log = getLogger('Featurizer')

__all__ = ['MDFeaturizer',
           'CustomFeature',
           ]


def _describe_atom(topology, index):
    """
    Returns a string describing the given atom

    :param topology:
    :param index:
    :return:
    """
    at = topology.atom(index)
    return "%s %i %s %i" % (at.residue.name, at.residue.index, at.name, at.index)


class CustomFeature(object):

    """
    A CustomFeature is the base class for all self defined features. You shall
    calculate your quantities in map method and return it as an ndarray.

    If you simply want to call a function with arguments, you do not need to derive.


    Parameters
    ----------
    func : function
        will be invoked with given args and kwargs on mapping traj
    args : list of positional args (optional)
    kwargs : named arguments

    Examples
    --------
    We use a FeatureReader to read md-trajectories and add a CustomFeature to
    transform coordinates:

    >>> reader = FeatureReader(...)
    >>> my_feature = CustomFeature(lambda x: 1.0 / x**2)
    >>> reader.featurizer.add_custom_feature(my_feature, output_dimension=3)

    Now a pipeline using this reader will apply the :math:`1 / x^2` transform on
    every frame being red.

    """

    def __init__(self, func=None, *args, **kwargs):
        self._func = func
        self._args = args
        self._kwargs = kwargs

    def describe(self):
        return "override me to get proper description!"

    def map(self, traj):
        feature = self._func(traj, self._args, self._kwargs)
        assert isinstance(feature, np.ndarray)
        return feature


class SelectionFeature:
    """
    Just provide the cartesian coordinates of a selection of atoms (could be simply all atoms).
    The coordinates are flattened as follows: [x1, y1, z1, x2, y2, z2, ...]

    """
    #TODO: Needs an orientation option

    def __init__(self, top, indexes):
        self.top = top
        self.indexes = np.array(indexes)
        self.prefix_label = "ATOM:"

    def describe(self):
        labels = []
        for i in self.indexes:
            labels.append("%s%s" % (self.prefix_label, _describe_atom(self.top, i)))
        return labels

    def map(self, traj):
        newshape = (traj.xyz.shape[0], 3*self.indexes.shape[0])
        return np.reshape(traj.xyz[:,self.indexes,:], newshape)


class DistanceFeature:

    def __init__(self, top, distance_indexes):
        self.top = top
        self.distance_indexes = np.array(distance_indexes)
        self.prefix_label = "DIST:"

    def describe(self):
        labels = []
        for pair in self.distance_indexes:
            labels.append("%s %s - %s" % (self.prefix_label,
                                          _describe_atom(self.top, pair[0]),
                                          _describe_atom(self.top, pair[1])))
        return labels

    def map(self, traj):
        return mdtraj.compute_distances(traj, self.distance_indexes)


class InverseDistanceFeature(DistanceFeature):

    def __init__(self, top, distance_indexes):
        DistanceFeature.__init__(self, top, distance_indexes)
        self.prefix_label = "INVDIST:"

    def map(self, traj):
        return 1.0 / mdtraj.compute_distances(traj, self.distance_indexes)


class ContactFeature(DistanceFeature):

    def __init__(self, top, distance_indexes, threshold):
        DistanceFeature.__init__(self, top, distance_indexes)
        self.prefix_label = "CONTACT:"
        self.threshold = threshold

    def map(self, traj):
        dists = mdtraj.compute_distances(traj, self.distance_indexes)
        res = np.zeros((len(traj), self.distance_indexes.shape[0]), dtype=np.float32)
        I = np.argwhere(dists > self.threshold)[0]
        res[I] = 1.0


class AngleFeature:<commit_msg>Documented Featurizer<commit_after>from pyemma.util.annotators import deprecated
__author__ = 'noe'

import mdtraj
from mdtraj.geometry.dihedral import _get_indices_phi, \
    _get_indices_psi, compute_dihedrals

import numpy as np
import warnings

from pyemma.util.log import getLogger

log = getLogger('Featurizer')

__all__ = ['MDFeaturizer',
           'CustomFeature',
           ]


def _describe_atom(topology, index):
    """
    Returns a string describing the given atom

    :param topology:
    :param index:
    :return:
    """
    at = topology.atom(index)
    return "%s %i %s %i" % (at.residue.name, at.residue.index, at.name, at.index)


class CustomFeature(object):

    """
    A CustomFeature is the base class for all self defined features. You shall
    calculate your quantities in map method and return it as an ndarray.

    If you simply want to call a function with arguments, you do not need to derive.


    Parameters
    ----------
    func : function
        will be invoked with given args and kwargs on mapping traj
    args : list of positional args (optional)
    kwargs : named arguments

    Examples
    --------
    We use a FeatureReader to read md-trajectories and add a CustomFeature to
    transform coordinates:

    >>> reader = FeatureReader(...)
    >>> my_feature = CustomFeature(lambda x: 1.0 / x**2)
    >>> reader.featurizer.add_custom_feature(my_feature, output_dimension=3)

    Now a pipeline using this reader will apply the :math:`1 / x^2` transform on
    every frame being red.

    """

    def __init__(self, func=None, *args, **kwargs):
        self._func = func
        self._args = args
        self._kwargs = kwargs

    def describe(self):
        return "override me to get proper description!"

    def map(self, traj):
        feature = self._func(traj, self._args, self._kwargs)
        assert isinstance(feature, np.ndarray)
        return feature


class SelectionFeature:
    """
    Just provide the cartesian coordinates of a selection of atoms (could be simply all atoms).
    The coordinates are flattened as follows: [x1, y1, z1, x2, y2, z2, ...]

    """
    #TODO: Needs an orientation option

    def __init__(self, top, indexes):
        self.top = top
        self.indexes = np.array(indexes)
        self.prefix_label = "ATOM:"

    def describe(self):
        labels = []
        for i in self.indexes:
            labels.append("%s%s" % (self.prefix_label, _describe_atom(self.top, i)))
        return labels

    def map(self, traj):
        newshape = (traj.xyz.shape[0], 3*self.indexes.shape[0])
        return np.reshape(traj.xyz[:,self.indexes,:], newshape)


class DistanceFeature:

    def __init__(self, top, distance_indexes):
        self.top = top
        self.distance_indexes = np.array(distance_indexes)
        self.prefix_label = "DIST:"

    def describe(self):
        labels = []
        for pair in self.distance_indexes:
            labels.append("%s %s - %s" % (self.prefix_label,
                                          _describe_atom(self.top, pair[0]),
                                          _describe_atom(self.top, pair[1])))
        return labels

    def map(self, traj):
        return mdtraj.compute_distances(traj, self.distance_indexes)


class InverseDistanceFeature(DistanceFeature):

    def __init__(self, top, distance_indexes):
        DistanceFeature.__init__(self, top, distance_indexes)
        self.prefix_label = "INVDIST:"

    def map(self, traj):
        return 1.0 / mdtraj.compute_distances(traj, self.distance_indexes)


class ContactFeature(DistanceFeature):

    def __init__(self, top, distance_indexes, threshold):
        DistanceFeature.__init__(self, top, distance_indexes)
        self.prefix_label = "CONTACT:"
        self.threshold = threshold

    def map(self, traj):
        dists = mdtraj.compute_distances(traj, self.distance_indexes)
        res = np.zeros((len(traj), self.distance_indexes.shape[0]), dtype=np.float32)
        I = np.argwhere(dists > self.threshold)
        res[I[:,0],I[:,1]] = 1.0
        return res


class AngleFeature:
{'fraction_non_alphanumeric': 0.1047280744732974, 'fraction_numerical': 0.006246937775600196, 'mean_word_length': 2.635351736420303, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5431959', 'n_tokens_mistral': 2509, 'n_tokens_neox': 2406, 'n_words': 690}