add data type; refactor methods

dpdata/bond_order_system.py

@@ -1,15 +1,13 @@
 #%%
 # Bond Order System
-from dpdata.system import System, LabeledSystem, check_System, load_format
+import numpy as np
+from dpdata.system import System, LabeledSystem, load_format, DataType, Axis
 import dpdata.rdkit.utils
 from dpdata.rdkit.sanitize import Sanitizer, SanitizeError
 from copy import deepcopy
 from rdkit.Chem import Conformer
 # import dpdata.rdkit.mol2
 
-def check_BondOrderSystem(data):
-    check_System(data)
-    assert ('bonds' in data.keys())
 
 class BondOrderSystem(System):
     '''
@@ -23,6 +21,11 @@ class BondOrderSystem(System):
                                         1 - single bond, 2 - double bond, 3 - triple bond, 1.5 - aromatic bond
         - `d_example['formal_charges']` : a numpy array of size 5 x 1
     '''
+    DTYPES = System.DTYPES + (
+        DataType("bonds", np.ndarray, (Axis.NBONDS, 3)),
+        DataType("formal_charges", np.ndarray, (Axis.NATOMS,)),
+    )
+
     def __init__(self,
                  file_name = None,
                  fmt = 'auto',
@@ -86,6 +89,7 @@ def __init__(self,
 
         if type_map:
             self.apply_type_map(type_map)
+        self.check_data()
 
     def from_fmt_obj(self, fmtobj, file_name, **kwargs):
         mol = fmtobj.from_bond_order_system(file_name, **kwargs)
@@ -104,9 +108,6 @@ def to_fmt_obj(self, fmtobj, *args, **kwargs):
             self.rdkit_mol.AddConformer(conf, assignId=True)
         return fmtobj.to_bond_order_system(self.data, self.rdkit_mol, *args, **kwargs)
 
-    def __repr__(self):
-        return self.__str__()
-
     def __str__(self):
         '''
             A brief summary of the system

dpdata/cp2k/output.py

@@ -268,7 +268,7 @@ def handle_single_xyz_frame(self, lines):
         #info_dict['atom_types'] = np.asarray(atom_types_list)
         info_dict['coords'] = np.asarray([coords_list]).astype('float32')
         info_dict['energies'] = np.array([energy]).astype('float32')
-        info_dict['orig']=[0,0,0]
+        info_dict['orig'] = np.zeros(3)
         return info_dict
 
 #%%

dpdata/plugins/ase.py

@@ -43,7 +43,7 @@ def from_system(self, atoms: "ase.Atoms", **kwargs) -> dict:
             'atom_types': atom_types,
             'cells': np.array([cells]).astype('float32'),
             'coords': np.array([coords]).astype('float32'),
-            'orig': [0,0,0],
+            'orig': np.zeros(3),
         }
         return info_dict
 

dpdata/pwmat/atomconfig.py

@@ -11,7 +11,7 @@ def _to_system_data_lower(lines) :
             for kk in range(idx+1,idx+1+3):
                 vector=[float(jj) for jj in lines[kk].split()[0:3]]
                 cell.append(vector)
-    system['cells'] = [np.array(cell)]
+    system['cells'] = np.array([cell])
     coord = []
     atomic_number = []
     atom_numbs = []
@@ -32,7 +32,7 @@ def _to_system_data_lower(lines) :
     for ii in np.unique(sorted(atomic_number)) :
         atom_numbs.append(atomic_number.count(ii))
     system['atom_numbs'] = [int(ii) for ii in atom_numbs]
-    system['coords'] = [np.array(coord)]
+    system['coords'] = np.array([coord])
     system['orig'] = np.zeros(3)
     atom_types = []
     for idx,ii in enumerate(system['atom_numbs']) :

dpdata/pymatgen/molecule.py

@@ -24,6 +24,6 @@ def to_system_data(file_name, protect_layer = 9) :
     # center = [c - h_cell_size for c in mol.center_of_mass]
     system['orig'] = np.array([0, 0, 0])
 
-    system['coords'] = [tmpcoord]
-    system['cells'] = [10.0 * np.eye(3)]
+    system['coords'] = np.array([tmpcoord])
+    system['cells'] = np.array([10.0 * np.eye(3)])
     return system