Source code for sknano.core.atoms._type_atoms
# -*- coding: utf-8 -*-
"""
===============================================================================
Atom classes with a 'type' attribute (:mod:`sknano.core.atoms._type_atoms`)
===============================================================================
An `Atom` class with a `type` attribute.
.. currentmodule:: sknano.core.atoms._type_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 sknano.core import dedupe
from ._atoms import Atom, Atoms
__all__ = ['TypeAtom', 'TypeAtoms']
@total_ordering
[docs]class TypeAtom(Atom):
"""An `Atom` class with an atom type attribute.
Parameters
----------
element : {str, int}, optional
A string representation of the element symbol or an integer specifying
an element atomic number.
type : int, optional
atom type
"""
def __init__(self, *args, type=1, **kwargs):
if 'atomtype' in kwargs:
type = kwargs['atomtype']
del kwargs['atomtype']
super().__init__(*args, **kwargs)
self.type = type
self.fmtstr = super().fmtstr + ", type={type!r}"
def __eq__(self, other):
return self.type == other.type and super().__eq__(other)
def __lt__(self, other):
return (self.type < other.type and super().__le__(other)) or \
(self.type <= other.type and super().__lt__(other))
def __dir__(self):
attrs = super().__dir__()
attrs.append('type')
return attrs
@property
def type(self):
""":attr:`~TypeAtom.type`."""
return self._type
@type.setter
def type(self, value):
"""Set :attr:`~TypeAtom.type`.
Parameters
----------
value : int
atom type
"""
if not isinstance(value, numbers.Number):
raise TypeError('Expected a number')
self._type = int(value)
@property
def atomtype(self):
return self.type
@atomtype.setter
def atomtype(self, value):
self.type = value
[docs] def todict(self):
super_dict = super().todict()
super_dict.update(dict(type=self.type))
return super_dict
[docs]class TypeAtoms(Atoms):
"""An `Atoms` sub-class for `TypeAtom`\ s.
A container class for `TypeAtom` objects.
Parameters
----------
atoms : {None, sequence, `TypeAtoms`}, optional
if not `None`, then a list of `TypeAtom` instance objects or an
existing `TypeAtoms` instance object.
"""
def __init__(self, atoms=None, **kwargs):
super().__init__(atoms, **kwargs)
self._typemap = {}
@property
def __atom_class__(self):
return TypeAtom
[docs] def sort(self, key=attrgetter('type'), reverse=False):
super().sort(key=key, reverse=reverse)
@property
def types(self):
""":class:`~numpy:numpy.ndarray` of :attr:`TypeAtom.type`\ s.
.. versionchanged:: 0.3.11
Returns an :class:`~numpy:numpy.ndarray` of
:attr:`TypeAtom.type`\ s instead of a :class:`~python:dict`
of type mapping. The :class:`~python:dict` of type
mappings is now assigned to the :attr:`~TypeAtoms.typemap`
attribute.
"""
return np.asarray([atom.type for atom in self])
@property
def atomtypes(self):
"""Alias for :attr:`~TypeAtoms.types`."""
return self.types
@property
def Ntypes(self):
"""Number of unique :attr:`~TypeAtoms.types`."""
return len(list(self.typemap.keys()))
@property
def typemap(self):
""":class:`python:dict` of :attr:`TypeAtom.type`\ s.
.. versionadded:: 0.3.11
"""
self._update_typemap()
return self._typemap
def _update_typemap(self):
[self.add_type(atom) for atom in self]
[docs] def add_type(self, atom):
"""Add atom type to :attr:`~TypeAtoms.typemap`.
Parameters
----------
atom : :class:`~sknano.core.atoms.TypeAtom`
A :class:`~sknano.core.atoms.TypeAtom` instance.
"""
self._typemap[atom.type] = {}
self._typemap[atom.type]['mass'] = atom.mass
[docs] def add_atomtype(self, atom):
"""Alias for :meth:`~TypeAtoms.add_type`."""
self.add_type(atom)
[docs] def add_types(self, atoms=None):
"""Add atom type for each atom in atoms to :attr:`TypeAtom.typemap` \
dictionary.
Parameters
----------
atoms : sequence
a list of `Atom` object instances
"""
try:
[self.add_type(atom) for atom in atoms]
except TypeError:
print('Expected an iterable sequence of `Atom` objects.')
[docs] def add_atomtypes(self, atoms=None):
"""Alias for :meth:`~TypeAtoms.add_types`."""
self.add_types(atoms=atoms)
[docs] def assign_unique_types(self, from_attr='element'):
"""Assign unique :attr:`TypeAtom.type`\s to each `TypeAtom` in \
`TypeAtoms` from an existing unique atom attribute.
.. versionchanged:: 0.3.11
Now accepts a keyword argument `from_attr`.
The assignment of unique :attr:`TypeAtom.type`\s is performed
by mapping an existing atom attribute (default: element)
to a unique integer, starting at 1.
Parameters
----------
from_attr : :class:`~python:str`
An existing atom attribute used to generate an attribute
mapping that maps the attribute to a unique atom
:attr:`~TypeAtom.type`.
"""
attrlist = [getattr(atom, from_attr) for atom in self]
attrmap = \
{attr: i for i, attr in enumerate(dedupe(attrlist), start=1)}
self.mapatomattr(from_attr, 'type', attrmap)
[docs] def get_types(self, asdict=False):
"""Return list of `TypeAtom` :attr:`TypeAtom.type`\ s.
Parameters
----------
asdict : :class:`python:bool`, optional
Returns
-------
:class:`python:list` if `asdict` is `False`
:class:`python:dict` if `asdict` is `True`
"""
if asdict:
return self.typemap
else:
return self.types.tolist()
[docs] def get_atomtypes(self, asdict=False):
"""Alias for :meth:`~TypeAtoms.get_types`."""
return self.get_types(asdict=asdict)