dpdata label system read ase db file

dpdata/ase/db.py

@@ -0,0 +1,41 @@
+from ase.db import connect
+import numpy as np
+
+
+def get_frames(fname, begin=0, step=1):
+    asedb = connect(fname)
+    num = asedb.count()
+    at0 = asedb.get(1).toatoms()
+
+    numbers = at0.numbers
+    sorted_numbers = numbers.argsort()
+    nat0 = at0[sorted_numbers]
+    chemical_symbols = nat0.get_chemical_symbols()
+    atom_names = list(set(chemical_symbols))
+    atom_numbs = [chemical_symbols.count(i) for i in atom_names]
+    atom_types = np.array([atom_names.index(i) for i in chemical_symbols])
+
+
+    all_coords = []
+    all_cells = []
+    all_energies = []
+    all_forces = []
+    all_virials = None
+
+    for i in range(begin, num, step):
+        ati = asedb.get(i+1)
+        ats_ = ati.toatoms()
+        numbers = ats_.numbers
+        sorted_numbers = numbers.argsort()
+        ats = ats_[sorted_numbers]
+        data = ati.data
+        energy = data['energy']
+        cell = ats.get_cell().view()
+        coord = ats.positions
+        forces = data['forces'][sorted_numbers]
+        all_coords.append(coord)
+        all_cells.append(cell)
+        all_forces.append(forces)
+        all_energies.append(energy)
+
+    return atom_names, atom_numbs, atom_types, np.array(all_cells), np.array(all_coords), np.array(all_energies), np.array(all_forces), all_virials

dpdata/system.py

@@ -8,6 +8,7 @@
 import dpdata.vasp.poscar
 import dpdata.vasp.xml
 import dpdata.vasp.outcar
+import dpdata.ase.db
 import dpdata.deepmd.raw
 import dpdata.deepmd.comp
 import dpdata.qe.traj
@@ -1137,6 +1138,21 @@ def from_vasp_outcar(self, file_name, begin = 0, step = 1) :
         # rotate the system to lammps convention
         self.rot_lower_triangular()
 
+    @register_from_funcs.register_funcs('db')
+    @register_from_funcs.register_funcs('ase/db')
+    def from_ase_db(self, file_name, begin = 0, step = 1) :
+        self.data['atom_names'], \
+            self.data['atom_numbs'], \
+            self.data['atom_types'], \
+            self.data['cells'], \
+            self.data['coords'], \
+            self.data['energies'], \
+            self.data['forces'], \
+            tmp_virial, \
+            = dpdata.ase.db.get_frames(file_name, begin = begin, step = step)
+
+        # rotate the system to lammps convention
+        self.rot_lower_triangular()
 
     def affine_map_fv(self, trans, f_idx) :
         assert(np.linalg.det(trans) != 0)