isolate get_atom_features

materialsvirtuallab · Dec 16, 2019 · 0a6db3f · 0a6db3f
1 parent 9da04cf
commit 0a6db3f
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/megnet/data/graph.py b/megnet/data/graph.py
@@ -100,7 +100,7 @@ def convert(self, structure: Structure, state_attributes: List = None) -> Dict:
         if np.size(np.unique(index1)) < len(atoms):
             raise RuntimeError("Isolated atoms found in the structure")
         else:
-            return {'atom': np.array(atoms, dtype='int32').tolist(),
+            return {'atom': atoms,
                     'bond': bonds,
                     'state': state_attributes,
                     'index1': index1,
@@ -118,7 +118,8 @@ def get_atom_features(structure) -> List[int]:
         Returns:
             List of atomic numbers
         """
-        return [i.specie.Z for i in structure]
+        return np.array([i.specie.Z for i in structure],
+                        dtype='int32').tolist()
 
     def __call__(self, structure: Structure) -> Dict:
         return self.convert(structure)
@@ -224,7 +225,7 @@ def convert(self, structure: Structure, state_attributes: List = None) -> Dict:
         if np.size(np.unique(index1)) < len(atoms):
             raise RuntimeError("Isolated atoms found in the structure")
         else:
-            return {'atom': np.array(atoms, dtype='int32').tolist(),
+            return {'atom': atoms,
                     'bond': bonds,
                     'state': state_attributes,
                     'index1': index1,