diff --git a/topology/core/element.py b/topology/core/element.py index dbd38c59c..4774e9634 100644 --- a/topology/core/element.py +++ b/topology/core/element.py @@ -2,34 +2,19 @@ from re import sub import numpy as np +from collections import namedtuple import unyt as u from topology.exceptions import TopologyError - -class Element(object): +class Element(namedtuple('Element', 'atomic_number, name, symbol, mass')): """An element.""" - def __init__(self, atomic_number=None, name=None, symbol=None, mass=None): - self._atomic_number = atomic_number - self._name = name - self._symbol = symbol - self._mass = mass - - @property - def atomic_number(self): - return self._atomic_number - - @property - def name(self): - return self._name - - @property - def symbol(self): - return self._symbol - - @property - def mass(self): - return self._mass + + def __str__(self): + return 'Element: {}, symbol: {}, atomic number: {}, mass: {}'.format( + self.name, self.symbol, + self.atomic_number, + self.mass) def element_by_symbol(symbol): diff --git a/topology/tests/base_test.py b/topology/tests/base_test.py index afc542083..4f17df524 100644 --- a/topology/tests/base_test.py +++ b/topology/tests/base_test.py @@ -4,7 +4,7 @@ from topology.core.box import Box from topology.core.topology import Topology -from topology.core.element import Element +from topology.core.element import Hydrogen from topology.core.site import Site from topology.core.atom_type import AtomType @@ -43,11 +43,11 @@ def topology_site(self): def _topology(sites=1): top = Topology() top.box = Box(lengths=[1, 1, 1]) - H = Element(name='H', symbol='H', mass=1) + H = Hydrogen site1 = Site(name='site1', element=H, atom_type=AtomType(name="at1", - mass=H.mass) + mass=H.mass), ) for i in range(sites): top.add_site(site1)