from rdkit.Chem import rdMolDescriptors
from ._base import Descriptor
from ._atomic_property import halogen
__all__ = (
'AtomCount',
)
[docs]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',)
@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 set(['H', 'Atom'])
def __str__(self):
return 'n' + self._type
[docs] 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