Merge pull request #109 from zezhong-zhang/devel

Sort atom according to the type_map in Multisystem

Author: amcadmus

dpdata/system.py

@@ -368,7 +368,8 @@ def sort_atom_names(self, type_map=None):
             # atom_names must be a subset of type_map
             assert (set(self.data['atom_names']).issubset(set(type_map)))
             # for the condition that type_map is a proper superset of atom_names
-            new_atoms = set(type_map) - set(self.data["atom_names"])
+            # new_atoms = set(type_map) - set(self.data["atom_names"])
+            new_atoms = [e for e in type_map if e not in self.data["atom_names"]]
             if new_atoms:
                 self.add_atom_names(new_atoms)
             # index that will sort an array by type_map
@@ -414,8 +415,8 @@ def formula(self):
         """
         Return the formula of this system, like C3H5O2
         """
-        return ''.join(["{}{}".format(symbol,numb) for symbol,numb in sorted(
-            zip(self.data['atom_names'], self.data['atom_numbs']))])
+        return ''.join(["{}{}".format(symbol,numb) for symbol,numb in 
+            zip(self.data['atom_names'], self.data['atom_numbs'])])
 
 
     def extend(self, systems):
@@ -1455,29 +1456,31 @@ def check_atom_names(self, system):
         """
         Make atom_names in all systems equal, prevent inconsistent atom_types.
         """
-        new_in_system = set(system["atom_names"]) - set(self.atom_names)
-        new_in_self = set(self.atom_names) - set(system["atom_names"])
+        # new_in_system = set(system["atom_names"]) - set(self.atom_names)
+        # new_in_self = set(self.atom_names) - set(system["atom_names"])
+        new_in_system = [e for e in system["atom_names"] if e not in self.atom_names]
+        new_in_self = [e for e in self.atom_names if e not in system["atom_names"]]
         if len(new_in_system):
             # A new atom_name appear, add to self.atom_names
             self.atom_names.extend(new_in_system)
-            self.atom_names.sort()
             # Add this atom_name to each system, and change their names
             new_systems = {}
             for each_system in self.systems.values():
                 each_system.add_atom_names(new_in_system)
-                each_system.sort_atom_names()
+                each_system.sort_atom_names(type_map=self.atom_names)
                 new_systems[each_system.formula] = each_system
             self.systems = new_systems
         if len(new_in_self):
             # Previous atom_name not in this system
             system.add_atom_names(new_in_self)
-        system.sort_atom_names()
+        system.sort_atom_names(type_map=self.atom_names)
 
     def from_quip_gap_xyz_file(self,file_name):
         # quip_gap_xyz_systems = QuipGapxyzSystems(file_name)
         # print(next(quip_gap_xyz_systems))
         for info_dict in QuipGapxyzSystems(file_name):
             system=LabeledSystem(data=info_dict)
+            system.sort_atom_names()
             self.append(system)