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)


[docs] 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` """ since = "1.0.0" __slots__ = ("_type", "_estate")
[docs] 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, version): 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"