Source code for mordred.EState

from enum import IntEnum
from functools import reduce

from rdkit.Chem import EState

from ._base import Descriptor
from ._util import parse_enum

try:
    import builtins
except ImportError:
    import __builtin__ as builtins


__all__ = ('AtomTypeEState',)


es_types = (
    'sLi', 'ssBe', 'ssssBe', 'ssBH', 'sssB', 'ssssB', 'sCH3', 'dCH2', 'ssCH2',
    'tCH', 'dsCH', 'aaCH', 'sssCH', 'ddC', 'tsC', 'dssC', 'aasC', 'aaaC',
    'ssssC', 'sNH3', 'sNH2', 'ssNH2', 'dNH', 'ssNH', 'aaNH', 'tN', 'sssNH',
    'dsN', 'aaN', 'sssN', 'ddsN', 'aasN', 'ssssN', 'sOH', 'dO', 'ssO', 'aaO',
    'sF', 'sSiH3', 'ssSiH2', 'sssSiH', 'ssssSi', 'sPH2', 'ssPH', 'sssP',
    'dsssP', 'sssssP', 'sSH', 'dS', 'ssS', 'aaS', 'dssS', 'ddssS', 'sCl',
    'sGeH3', 'ssGeH2', 'sssGeH', 'ssssGe', 'sAsH2', 'ssAsH', 'sssAs', 'sssdAs',
    'sssssAs', 'sSeH', 'dSe', 'ssSe', 'aaSe', 'dssSe', 'ddssSe', 'sBr', 'sSnH3',
    'ssSnH2', 'sssSnH', 'ssssSn', 'sI', 'sPbH3', 'ssPbH2', 'sssPbH', 'ssssPb'
)

es_type_set = set(es_types)


class EStateBase(Descriptor):
    __slots__ = ()
    explicit_hydrogens = False


class EStateCache(EStateBase):
    __slots__ = ()

    def parameters(self):
        return ()

    def calculate(self):
        return EState.TypeAtoms(self.mol), EState.EStateIndices(self.mol)


class AggrType(IntEnum):
    __slots__ = ()

    count = 1
    sum = 2
    max = 3
    min = 4

    @property
    def as_argument(self):
        return self.name


aggr_names = (
    (AggrType.count, 'N'),
    (AggrType.sum, 'S'),
    (AggrType.max, 'MAX'),
    (AggrType.min, 'MIN'),
)


aggr_name_dict = dict(aggr_names)


class AtomTypeEState(EStateBase):
    r"""atom type e-state descriptor.

    :type type: str
    :param type: one of aggr_types

    :type estate: str
    :param estate: one of es_types

    :returns: NaN when type in ['min', 'max'] and :math:`N_{\rm X} = 0`

    References
        * :doi:`10.1021/ci00028a014`
    """
    __slots__ = ('_type', '_estate',)

    aggr_types = tuple(a.name for a in AggrType)

    es_types = es_types

    @classmethod
    def preset(cls):
        return (
            cls(a, t)
            for a in AggrType
            for t in es_types
        )

    def __str__(self):
        aggr = aggr_name_dict[self._type]

        return aggr + self._estate

    def parameters(self):
        return self._type, self._estate

    def __init__(self, type='count', estate='sLi'):
        assert estate in es_type_set

        self._type = parse_enum(AggrType, type)
        self._estate = estate

    def dependencies(self):
        return {'E': EStateCache()}

    def calculate(self, E):
        if self._type == AggrType.count:
            return reduce(lambda a, b: a + b, E[0]).count(self._estate)

        indices = map(
            lambda e: e[1],
            filter(lambda e: self._estate in e[0], zip(*E))
        )

        with self.rethrow_na(ValueError):
            return getattr(builtins, self._type.name)(indices)

    @property
    def rtype(self):
        r"""
        * 'count': :py:class:`int`
        * 'other': :py:class:`float`
        """

        return int if self._type == AggrType.count else float

    _extra_docs = 'aggr_types', 'es_types'