From eedc02161a93af9b9b59e07c90b7ef47b52e19c6 Mon Sep 17 00:00:00 2001 From: Aaron Virshup Date: Fri, 30 Jun 2017 20:56:46 -0700 Subject: [PATCH] Atomic properties and trajectories play nice --- moldesign/molecules/properties.py | 6 ++---- moldesign/molecules/trajectory.py | 18 ++++++++++-------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/moldesign/molecules/properties.py b/moldesign/molecules/properties.py index c585a07..4e9c05e 100644 --- a/moldesign/molecules/properties.py +++ b/moldesign/molecules/properties.py @@ -16,9 +16,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import numpy as np -import collections - from . import toplevel from .. import utils, units @@ -107,7 +104,8 @@ def __setitem__(self, atom, value): def __getitem__(self, atom): return super().__getitem__(self._getkey(atom)) - def _getkey(self, atom): + @staticmethod + def _getkey(atom): if isinstance(atom, int): k = atom else: diff --git a/moldesign/molecules/trajectory.py b/moldesign/molecules/trajectory.py index f517356..1a1fbdc 100644 --- a/moldesign/molecules/trajectory.py +++ b/moldesign/molecules/trajectory.py @@ -24,7 +24,7 @@ import moldesign as mdt from .. import helpers, utils from .. import units as u -from .molecule import MolecularProperties +from . import MolecularProperties, AtomicProperties from . import toplevel @@ -105,17 +105,17 @@ def _arrayslice(self, attr): def __getattr__(self, item): if item in ('traj', 'index', 'real_atom'): - raise AttributeError('_TrajAtom.%s not assigned (pickle issue?)' % item) + raise AttributeError('Internal Error: _TrajAtom.%s has not been initialized' % item) - try: # try to get a time-dependent version - trajslice = getattr(self.traj, item) - except AttributeError: # not time-dependent - look for an atomic property + try: # first, look for time-dependent properties + val = getattr(self.traj, item) + except AttributeError: # otherwise, just return an atomic property return getattr(self.real_atom, item) - if trajslice[0]['type'] != 'atomic': + if not isinstance(val[0], AtomicProperties): raise ValueError('%s is not an atomic quantity' % item) else: - return u.array([f[self.real_atom] for f in trajslice]) + return u.array([f[self.real_atom] for f in val]) def __dir__(self): attrs = [self.ATOMIC_ARRAYS.get(x, x) for x in dir(self.traj)] @@ -346,7 +346,7 @@ def new_frame(self, properties=None, **additional_data): newpl.append(None) self.properties[key] = newpl - # TODO: less flubby way of keeping track of # of frames + # TODO: less awkward way of keeping track of # of frames self.frames.append(Frame(self, self.num_frames)) @@ -365,6 +365,8 @@ def _new_property(self, key, value): if self.num_frames != 0: proplist = [None] * self.num_frames proplist.append(value) + elif isinstance(value, (AtomicProperties, dict)): + proplist = [value] else: try: proplist = self.unit_system.convert(u.array([value]))