Atomic properties and trajectories play nice

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:

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]))