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

    def description(self):
        d = {
            self.count: "number",
            self.sum: "sum",
            self.max: "max",
            self.min: "min",
        }
        return d[self]


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",)

    def description(self):
        return "{} of {}".format(self._type.description(), self._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"""Return type.

        * "count": :py:class:`int`
        * "other": :py:class:`float`
        """
        return int if self._type == AggrType.count else float

    _extra_docs = "aggr_types", "es_types"