Updated atom table to support saving labels (if present)

Author: avmarchenko

exatomic/algorithms/diffusion.py

@@ -34,5 +34,5 @@ def einstein_relation(universe, length='cm', time='s'):
     msd = absolute_squared_displacement(universe).mean(axis=1)
     t = universe.frame['time'].values
     msd *= Length['au', length]**2
-    t /= Time['ps', time]
+    t /= Time['au', time]
     return msd / (6 * t)

exatomic/atom.py

@@ -12,7 +12,7 @@
 from traitlets import Dict, Unicode
 from exa.numerical import DataFrame, SparseDataFrame, Series
 from exa.relational.isotope import (symbol_to_color, symbol_to_radius,
-                                   symbol_to_element_mass)
+                                    symbol_to_element_mass)
 from exatomic.error import PeriodicUniverseError
 from exatomic.algorithms.distance import minimal_image_counts
 from exatomic.algorithms.geometry import make_small_molecule
@@ -51,7 +51,8 @@ class Atom(DataFrame):
     _precision = {'x': 2, 'y': 2, 'z': 2}
     _index = 'atom'
     _cardinal = ('frame', np.int64)
-    _categories = {'symbol': str, 'set': np.int64, 'molecule': np.int64}
+    _categories = {'symbol': str, 'set': np.int64, 'molecule': np.int64,
+                   'label': np.int64}
     _traits = ['x', 'y', 'z', 'set']
     _columns = ['x', 'y', 'z', 'symbol']
 

exatomic/container.py

@@ -113,6 +113,7 @@ def compute_atom_two(self, mapper=None, bond_extra=0.45):
             self.projected_atom = projected_atom
         else:
             self.atom_two = compute_atom_two(self, mapper, bond_extra)
+        self._traits_need_update = True
 
     def compute_bonds(self, mapper=None, bond_extra=0.45):
         """
@@ -123,6 +124,7 @@ def compute_bonds(self, mapper=None, bond_extra=0.45):
         """
         self.atom_two.compute_bonds(self.atom['symbol'], mapper=mapper, bond_extra=bond_extra)
         self.compute_molecule()
+        self._traits_need_update = True
 
     def compute_bond_count(self):
         """
@@ -175,6 +177,7 @@ def add_field(self, field):
                 self.field = AtomicField(new_field, field_values=new_field_values)
         else:
             raise TypeError('field must be an instance of exatomic.field.AtomicField or a list of them')
+        self._traits_need_update = True
 
     def add_molecular_orbitals(self, *field_params, mocoefs=None, vector=None):
         """
@@ -190,6 +193,7 @@ def add_molecular_orbitals(self, *field_params, mocoefs=None, vector=None):
         if not hasattr(self, '_basis_set_order'):
             print('Warning: without the basis_set_order, MOs are likely incorrect.')
         _add_mos_to_universe(self, *field_params, mocoefs=mocoefs, vector=vector)
+        self._traits_need_update = True
 
     def update_molecular_orbitals(self, *field_params, mocoefs=None, vector=None):
         """
@@ -203,6 +207,7 @@ def update_molecular_orbitals(self, *field_params, mocoefs=None, vector=None):
         if not hasattr(self, 'basis_functions'):
             raise AttributeError('Universe has no basis functions, add_molecular_orbitals first')
         _update_mos(self, *field_params, mocoefs=mocoefs, vector=vector)
+        self._traits_need_update = True
 
     def _custom_traits(self):
         """