qe/cp/traj: load cell data if there is any. check the consistency of time

dpdata/qe/traj.py

@@ -98,15 +98,16 @@ def _load_pos_block(fp, natoms) :
     head = fp.readline()
     if not head:
         # print('get None')
-        return None
+        return None, None
     else :
+        ss = head.split()[0]
         blk = []
         for ii in range(natoms) :
             newline = fp.readline()
             if not newline :
-                return None
+                return None, None
             blk.append([float(jj) for jj in newline.split()])
-        return blk
+        return blk, ss
 
 
 def load_data(fname, 
@@ -115,18 +116,20 @@ def load_data(fname,
               step = 1,
               convert = 1.) :
     coords = []
+    steps = []
     cc = 0
     with open(fname) as fp:
         while True:
-            blk = _load_pos_block(fp, natoms)
+            blk, ss = _load_pos_block(fp, natoms)
             if blk == None :
                 break
             else :
                 if cc >= begin and (cc - begin) % step == 0 :
                     coords.append(blk)
+                    steps.append(ss)
             cc += 1
     coords= convert * np.array(coords)
-    return coords
+    return coords, steps
 
 
 # def load_pos(fname, natoms) :
@@ -145,14 +148,17 @@ def load_data(fname,
 
 def load_energy(fname, begin = 0, step = 1) :
     data = np.loadtxt(fname)
+    steps = []
+    for ii in data[begin::step,0]:
+        steps.append('%d'%ii)
     with open(fname) as fp:
         line = fp.readline()
         if line :
             nw = len(line.split())
         else :
             return None
     data = np.reshape(data, [-1, nw])
-    return energy_convert * data[begin::step,5]
+    return energy_convert * data[begin::step,5], steps
 
 
 # def load_force(fname, natoms) :
@@ -176,35 +182,48 @@ def to_system_data(input_name, prefix, begin = 0, step = 1) :
         data['atom_types'], \
         cell \
         = load_param_file(input_name)
-    data['coords'] \
+    data['coords'], csteps\
         = load_data(prefix + '.pos', 
                     np.sum(data['atom_numbs']), 
                     begin = begin, 
                     step = step, 
                     convert = length_convert)
     data['orig'] = np.zeros(3)
-    data['cells'] = np.tile(cell, (data['coords'].shape[0], 1, 1))
-    return data
+    try :
+        data['cells'], tmp_steps \
+            = load_data(prefix + '.cel', 
+                        3, 
+                        begin = begin, 
+                        step = step, 
+                        convert = length_convert)
+        assert(csteps == tmp_steps), "the step key between files are not consistent"
+    except FileNotFoundError :
+        data['cells'] = np.tile(cell, (data['coords'].shape[0], 1, 1))
+    return data, csteps
 
 
 def to_system_label(input_name, prefix, begin = 0, step = 1) :
     atom_names, atom_numbs, atom_types, cell = load_param_file(input_name)
-    energy = load_energy(prefix + '.evp', 
-                         begin = begin, 
-                         step = step)
-    force = load_data(prefix + '.for', 
-                      np.sum(atom_numbs), 
-                      begin = begin,
-                      step = step, 
-                      convert = force_convert)
-    return energy, force
+    energy, esteps = load_energy(prefix + '.evp', 
+                                 begin = begin, 
+                                 step = step)
+    force, fsteps = load_data(prefix + '.for', 
+                              np.sum(atom_numbs), 
+                              begin = begin,
+                              step = step, 
+                              convert = force_convert)
+    assert(esteps == fsteps), "the step key between files are not consistent "
+    return energy, force, esteps
 
 
 if __name__ == '__main__':
-    atom_names, atom_numbs, atom_types, cell = load_param_file('oh-md.in')
-    coords = load_pos('oh-md.pos', np.sum(atom_numbs))
+    prefix='nacl'
+    atom_names, atom_numbs, atom_types, cell = load_param_file(prefix+'.in')
+    coords = load_data(prefix+'.pos', np.sum(atom_numbs))
+    cells = load_data(prefix+'.cel', 3)
     print(atom_names)
     print(atom_numbs)
     print(atom_types)
-    print(cell)
+    print(cells)
     print(coords.shape)
+    print(cells.shape)

dpdata/system.py

@@ -525,7 +525,7 @@ def from_qe_cp_traj(self,
                         prefix,
                         begin = 0,
                         step = 1) :
-        self.data = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step)
+        self.data, _ = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step)
         self.data['coords'] \
             = dpdata.md.pbc.apply_pbc(self.data['coords'],
                                       self.data['cells'],
@@ -991,13 +991,14 @@ def from_deepmd_raw(self, folder, type_map = None) :
 
 
     def from_qe_cp_traj(self, prefix, begin = 0, step = 1) :
-        self.data = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step)
+        self.data, cs = dpdata.qe.traj.to_system_data(prefix + '.in', prefix, begin = begin, step = step)
         self.data['coords'] \
             = dpdata.md.pbc.apply_pbc(self.data['coords'],
                                       self.data['cells'],
             )
-        self.data['energies'], self.data['forces'] \
+        self.data['energies'], self.data['forces'], es \
             = dpdata.qe.traj.to_system_label(prefix + '.in', prefix, begin = begin, step = step)
+        assert(cs == es), "the step key between files are not consistent" 
         self.rot_lower_triangular()
 
     def from_qe_pw_scf(self, file_name) :

tests/qe.traj/traj6.cel

@@ -0,0 +1,24 @@
+195
+5.157620138324029213e+00 0 0
+8.882263018337732685e-02 5.782217839848955876e+00 0
+4.531494312205264219e-02 1.857154025688012577e-01 5.855553673164934025e+00
+200
+5.359985500701728967e+00 0 0
+3.585941820098031974e-01 5.317218997480877896e+00 0
+7.606780476053129902e-01 7.811107228901693622e-01 5.715864930517207121e+00
+201
+5.235484852416917079e+00 0 0
+5.612267823951659906e-01 5.640178212424978632e+00 0
+1.561235654197479228e-02 2.819721738692628765e-01 5.868536463888631260e+00
+202
+5.735200354152383717e+00 0 0
+8.818112752306928037e-02 5.003946763988647461e+00 0
+5.346029912365874992e-01 4.453576669882056693e-01 5.311832768820492490e+00
+203
+5.969007839117177916e+00 0 0
+7.974158979687617776e-02 5.086270090760550922e+00 0
+2.286579398249665163e-01 1.449088120223311904e-01 5.588188387987865546e+00
+204
+5.308510801020571712e+00 0 0
+3.076052782312116429e-01 5.279388982187173340e+00 0
+4.321921336152507731e-01 8.121110815096156399e-01 5.301664983741235737e+00

tests/qe.traj/traj6.evp

@@ -1,6 +1,6 @@
     195  9.433649E-03  1.290500E-05  0.000000E+00  1.057182E-02      -1100.03076389      -1100.03076389      -1100.03075434      -1100.03073209  1.300606E+04      -3.31610
     200  9.675537E-03  7.090171E-05  0.000000E+00  1.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
-    200  9.675537E-03  7.090171E-05  0.000000E+00  2.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
-    200  9.675537E-03  7.090171E-05  0.000000E+00  3.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
-    200  9.675537E-03  7.090171E-05  0.000000E+00  4.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
-    200  9.675537E-03  7.090171E-05  0.000000E+00  5.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
+    201  9.675537E-03  7.090171E-05  0.000000E+00  2.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
+    202  9.675537E-03  7.090171E-05  0.000000E+00  3.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
+    203  9.675537E-03  7.090171E-05  0.000000E+00  4.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683
+    204  9.675537E-03  7.090171E-05  0.000000E+00  5.257319E+00      -1100.03189493      -1100.03189493      -1100.03076015      -1100.03068120  1.300606E+04      -3.31683

tests/qe.traj/traj6.for

@@ -4,15 +4,15 @@
     200  0.00967554
      0.13452693616244E-01    -0.68546285592109E-02    -0.28341150176630E-02
      0.23859341498989E-01    -0.15249791118496E-01    -0.77171937738584E-02
-    200  0.00967554
+ 201  0.00967554
      2.13452693616244E-01    -0.68546285592109E-02    -0.28341150176630E-02
      0.23859341498989E-01    -0.15249791118496E-01    -0.77171937738584E-02
-    200  0.00967554
+ 202  0.00967554
      3.13452693616244E-01    -0.68546285592109E-02    -0.28341150176630E-02
      0.23859341498989E-01    -0.15249791118496E-01    -0.77171937738584E-02
-    200  0.00967554
+ 203  0.00967554
      4.13452693616244E-01    -0.68546285592109E-02    -0.28341150176630E-02
      0.23859341498989E-01    -0.15249791118496E-01    -0.77171937738584E-02
-    200  0.00967554
+ 204  0.00967554
      5.13452693616244E-01    -0.68546285592109E-02    -0.28341150176630E-02
      0.23859341498989E-01    -0.15249791118496E-01    -0.77171937738584E-02

tests/qe.traj/traj6.pos

@@ -4,15 +4,15 @@
     200  0.00967554
      0.32121855340601E+02     0.17164728967776E+02    -0.36182628967521E+01
      0.14321036136090E+01     0.54322965963498E+01     0.16038534464162E+02
-    200  0.00967554
+    201  0.00967554
      2.32121855340601E+02     0.17164728967776E+02    -0.36182628967521E+01
      0.14321036136090E+01     0.54322965963498E+01     0.16038534464162E+02
-    200  0.00967554
+    202  0.00967554
      3.32121855340601E+02     0.17164728967776E+02    -0.36182628967521E+01
      0.14321036136090E+01     0.54322965963498E+01     0.16038534464162E+02
-    200  0.00967554
+    203  0.00967554
      4.32121855340601E+02     0.17164728967776E+02    -0.36182628967521E+01
      0.14321036136090E+01     0.54322965963498E+01     0.16038534464162E+02
-    200  0.00967554
+    204  0.00967554
      5.32121855340601E+02     0.17164728967776E+02    -0.36182628967521E+01
      0.14321036136090E+01     0.54322965963498E+01     0.16038534464162E+02

tests/test_qe_cp_traj_skipload.py

@@ -38,3 +38,21 @@ def setUp(self):
         self.f_places = 6
         self.v_places = 4
 
+    def test_cell(self):
+        ref_cell = [5.359985500701728967e+00, 0, 0,
+                    3.585941820098031974e-01, 5.317218997480877896e+00, 0,
+                    7.606780476053129902e-01, 7.811107228901693622e-01, 5.715864930517207121e+00 ]
+        ref_cell = 0.52917721067 * np.array(ref_cell).reshape(3,3)
+
+        for ii in range(3) :
+            for jj in range(3) :
+                self.assertEqual(self.system_1.data['cells'][0][ii][jj], ref_cell[ii][jj])
+
+        ref_cell = [5.308510801020571712e+00, 0, 0,
+                    3.076052782312116429e-01, 5.279388982187173340e+00, 0,
+                    4.321921336152507731e-01, 8.121110815096156399e-01, 5.301664983741235737e+00]
+        ref_cell = 0.52917721067 * np.array(ref_cell).reshape(3,3)
+
+        for ii in range(3) :
+            for jj in range(3) :
+                self.assertEqual(self.system_1.data['cells'][-1][ii][jj], ref_cell[ii][jj])