Source code for sknano.core.atoms._id_atoms

# -*- coding: utf-8 -*-
"""
===============================================================================
Atom classes with id attributes (:mod:`sknano.core.atoms._id_atoms`)
===============================================================================

.. currentmodule:: sknano.core.atoms._id_atoms

"""
from __future__ import absolute_import, division, print_function
from __future__ import unicode_literals

__docformat__ = 'restructuredtext en'

from functools import total_ordering
from operator import attrgetter
import numbers

import numpy as np

from ._atoms import Atom, Atoms

__all__ = ['IDAtom', 'IDAtoms']


@total_ordering
[docs]class IDAtom(Atom): """An `Atom` class with id attributes. Parameters ---------- element : {str, int}, optional A string representation of the element symbol or an integer specifying an element atomic number. id : int, optional atom ID mol : int, optional molecule ID """ def __init__(self, *args, id=0, mol=0, **kwargs): if 'atomID' in kwargs: id = kwargs['atomID'] del kwargs['atomID'] if 'moleculeID' in kwargs: mol = kwargs['moleculeID'] del kwargs['moleculeID'] super().__init__(*args, **kwargs) self.id = id self.mol = mol self.fmtstr = super().fmtstr + ", id={id!r}, mol={mol!r}" def __eq__(self, other): return self.id == other.id and self.mol == other.mol and \ super().__eq__(other) def __lt__(self, other): return ((self.id < other.id and self.mol <= other.mol and super().__le__(other)) or (self.id <= other.id and self.mol < other.mol and super().__le__(other)) or (self.id <= other.id and self.mol <= other.mol and super().__lt__(other))) def __dir__(self): attrs = super().__dir__() attrs.extend(['id', 'mol']) return attrs @property def id(self): """:attr:`~IDAtom.id`.""" return self._id @id.setter def id(self, value): """Set atom id. Parameters ---------- value : int atom ID """ if not isinstance(value, numbers.Number): raise TypeError('Expected a number') self._id = int(value) @property def mol(self): """:attr:`~IDAtom.mol`.""" return self._mol @mol.setter def mol(self, value): """Set :attr:`~IDAtom.mol`. Parameters ---------- value : int molecule ID """ if not isinstance(value, numbers.Number): raise TypeError('Expected a number') self._mol = int(value) @property def atomID(self): return self.id @atomID.setter def atomID(self, value): self.id = value @property def moleculeID(self): return self.mol @moleculeID.setter def moleculeID(self, value): self.mol = value
[docs] def todict(self): super_dict = super().todict() super_dict.update(dict(id=self.id, mol=self.mol)) return super_dict
[docs]class IDAtoms(Atoms): """An `Atoms` sub-class for `IDAtom`\ s. Sub-class of `Atoms` class, and a container class for lists of :class:`~sknano.core.atoms.IDAtom` instances. Parameters ---------- atoms : {None, sequence, `IDAtoms`}, optional if not `None`, then a list of `IDAtom` instance objects or an existing `IDAtoms` instance object. """ @property def __atom_class__(self): return IDAtom
[docs] def sort(self, key=attrgetter('mol', 'id'), reverse=False): super().sort(key=key, reverse=reverse)
@property def ids(self): """Return array of `IDAtom.id`\ s.""" if len(set([atom.id for atom in self])) != self.Natoms: self.assign_unique_ids() return np.asarray([atom.id for atom in self]) @property def mols(self): """Return array of `IDAtom.mol`\ s.""" return np.asarray([atom.mol for atom in self]) @property def atom_ids(self): """Alias for :attr:`IDAtoms.ids`.""" return self.ids @property def mol_ids(self): """Alias for :attr:`IDAtoms.mols`.""" return self.mols @property def molecule_ids(self): """Alias for :attr:`IDAtoms.mols`.""" return self.mols
[docs] def assign_unique_ids(self, starting_id=1): """Assign unique :attr:`IDAtom.id` to each `IDAtom` in `IDAtoms`.""" [setattr(atom, 'id', i) for i, atom in enumerate(self, start=starting_id)]
[docs] def filter_ids(self, atom_ids, invert=False): """Filter `Atoms` by :attr:`IDAtoms.ids` in `atom_ids`. .. versionchanged:: 0.3.11 Filters `Atoms` **in-place**. Use :meth:`~IDAtoms.filtered_ids` to get a new list of `Atoms`. Parameters ---------- atom_ids : array_like invert : bool, optional """ mask = np.in1d(self.ids, atom_ids, invert=invert).nonzero() self.data = np.asarray(self)[mask].tolist()
[docs] def filtered_ids(self, atom_ids, invert=False): """Return new `Atoms` filtered by :attr:`IDAtoms.ids` in `atom_ids`. .. versionadded:: 0.3.11 Parameters ---------- atom_ids : array_like invert : bool, optional Returns ------- filtered_atoms : `Atoms` An instance of `Atoms` (sub)class. """ mask = np.in1d(self.ids, atom_ids, invert=invert).nonzero() return self.__class__(atoms=np.asarray(self)[mask].tolist(), **self.kwargs)
[docs] def get_atom(self, id): """Get `IDAtom` with :attr:`Xatom.id` == `id`. Parameters ---------- id : int Returns ------- atom : `IDAtom` or `None` `IDAtom` instance if `IDAtoms` contains `IDAtom` with :attr:`IDAtom.id` == `id`, otherwise `None` """ try: return self[np.where(self.ids == id)[0]] except TypeError: print('No atom with id = {}'.format(id)) return None