Source code for mordred.AtomCount

from rdkit.Chem import rdMolDescriptors

from ._base import Descriptor
from ._atomic_property import halogen

__all__ = (
    "AtomCount",
)


_desc_conv = {
    "Atom": "all",
    "Bridgehead": "bridgehead",
    "HeavyAtom": "heavy",
    "Spiro": "spiro",
    "X": "halogen",
}


class AtomCount(Descriptor):
    r"""atom count descriptor.

    :type type: str
    :param type: type to count.

        * "Atom"
        * "HeavyAtom"
        * "Spiro"
        * "Bridgehead"
        * "X" - all halogen
        * element symbol
    """

    __slots__ = ("_type",)

    def description(self):
        return "number of {} atoms".format(_desc_conv.get(self._type, self._type))

    @classmethod
    def preset(cls):
        return map(cls, [
            "Atom", "HeavyAtom", "Spiro", "Bridgehead",
            "H", "B", "C", "N", "O", "S", "P", "F", "Cl", "Br", "I", "X",
        ])

    @property
    def explicit_hydrogens(self):
        u"""Require explicit_hydrogens when type is "H" or "Atom"."""
        return self._type in {"H", "Atom"}

    def __str__(self):
        return "n" + self._type

    def parameters(self):
        return self._type,

    def __init__(self, type="Atom"):
        self._type = type

    def _calc_X(self):
        X = halogen
        return sum(a.GetAtomicNum() in X for a in self.mol.GetAtoms())

    def _calc(self):
        return sum(a.GetSymbol() == self._type for a in self.mol.GetAtoms())

    def _calc_all(self):
        return self.mol.GetNumAtoms()

    def calculate(self):
        if self._type == "X":
            return self._calc_X()
        elif self._type in ["Atom", "HeavyAtom"]:
            return self._calc_all()
        elif self._type == "Spiro":
            return rdMolDescriptors.CalcNumSpiroAtoms(self.mol)
        elif self._type == "Bridgehead":
            return rdMolDescriptors.CalcNumBridgeheadAtoms(self.mol)
        else:
            return self._calc()

    rtype = int