fix bug in reading trajs with random tid (#377)

fix #376

Co-authored-by: Sikai Yao <[email protected]>

Author: MabinogiX

dpdata/lammps/dump.py

@@ -31,9 +31,11 @@ def get_atype(lines, type_idx_zero = False) :
     tidx = keys.index('type') - 2
     atype = []
     for ii in blk :
-        atype.append([int(ii.split()[id_idx]), int(ii.split()[tidx])])
+        atype.append([int(ii.split()[tidx]), int(ii.split()[id_idx])])
+    # sort with type id
     atype.sort()
     atype = np.array(atype, dtype = int)    
+    atype = atype[:, ::-1]
     if type_idx_zero :
         return atype[:,1] - 1
     else :
@@ -76,16 +78,17 @@ def safe_get_posi(lines,cell,orig=np.zeros(3), unwrap=False) :
     assert coord_tp_and_sf is not None, 'Dump file does not contain atomic coordinates!'
     coordtype, sf, uw = coord_tp_and_sf
     id_idx = keys.index('id') - 2
+    tidx = keys.index('type') - 2
     xidx = keys.index(coordtype[0])-2
     yidx = keys.index(coordtype[1])-2
     zidx = keys.index(coordtype[2])-2
     sel = (xidx, yidx, zidx)
     posis = []
     for ii in blk :
         words = ii.split()
-        posis.append([float(words[id_idx]), float(words[xidx]), float(words[yidx]), float(words[zidx])])
+        posis.append([float(words[tidx]), float(words[id_idx]), float(words[xidx]), float(words[yidx]), float(words[zidx])])
     posis.sort()
-    posis = np.array(posis)[:,1:4]
+    posis = np.array(posis)[:,2:5]
     if not sf:
         posis = (posis - orig) @ np.linalg.inv(cell)  # Convert to scaled coordinates for unscaled coordinates
     if uw and unwrap:

tests/lammps/traj_with_random_type_id.dump

@@ -0,0 +1,50 @@
+ITEM: TIMESTEP
+0
+ITEM: NUMBER OF ATOMS
+16
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+0.0000000000000000e+00 6.8043376809999998e+00 2.5385198599999999e-02
+0.0000000000000000e+00 6.7821075796999999e+00 1.8630761460000000e-01
+0.0000000000000000e+00 6.6801861338000004e+00 6.5204177000000002e-02
+ITEM: ATOMS id type x y z
+1 4 3.48873 0.0697213 6.67774
+2 4 3.38621 0.033338 3.34239
+3 4 1.79424 1.7281 5.01015
+4 4 3.48973 3.42896 6.67795
+5 4 3.40064 3.39148 3.34188
+6 2 6.69832 3.39136 3.34005
+7 3 1.80744 5.08708 5.01099
+8 3 5.10512 5.08007 5.01272
+9 3 1.70086 1.69544 1.66979
+10 4 5.09069 1.72876 5.00917
+11 4 0.119885 6.74841 3.33869
+12 4 4.99379 1.69262 1.67183
+13 4 0.199838 3.4185 6.67565
+14 4 1.7213 5.05235 1.66373
+15 4 0.21494 6.77616 6.67623
+16 6 5.00691 5.05 1.66532
+ITEM: TIMESTEP
+10
+ITEM: NUMBER OF ATOMS
+16
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+3.0951719137647604e-02 6.7713982144168243e+00 2.5146837349522749e-02
+3.1535098850918430e-02 6.7499602284333751e+00 1.8455822840494820e-01
+3.1362715442244227e-02 6.6488234183577575e+00 6.4591924584292706e-02
+ITEM: ATOMS id type x y z
+1 6 3.52465 0.174767 6.64949
+2 4 3.44881 6.57204 3.4593
+3 3 1.85127 1.64364 4.94466
+4 4 3.4242 3.50712 6.53701
+5 4 3.46382 3.47183 3.41008
+6 4 6.63593 3.49936 3.46086
+7 4 1.85117 5.11268 4.96295
+8 3 5.18664 4.9445 5.05251
+9 3 1.70736 1.6591 1.6432
+10 4 5.23527 1.84126 5.03198
+11 4 6.59374 6.65872 3.38238
+12 4 5.00141 1.8369 1.76495
+13 4 6.64861 3.4664 6.51584
+14 4 1.60995 4.92415 1.66769
+15 4 6.62681 0.172193 6.54904
+16 2 4.88678 5.15478 1.589

tests/test_lammps_read_from_trajs.py

@@ -0,0 +1,44 @@
+import os
+import numpy as np
+import unittest
+from context import dpdata
+
+class TestLmpReadFromTrajsWithRandomTypeId(unittest.TestCase):
+    
+    def setUp(self): 
+        self.system = \
+            dpdata.System(os.path.join('lammps', 'traj_with_random_type_id.dump'), fmt = 'lammps/dump', type_map = ["Ta","Nb","W","Mo","V","Al"])
+    
+    def test_nframes (self) :
+        atype = self.system['atom_types'].tolist()
+        self.assertTrue(atype == [1, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 5])
+        
+        coord = self.system['coords'].reshape([2, -1])
+
+        coord0_std = np.array([6.69832  , 3.39136  , 3.34005  , 1.80744  , 5.08708  , 5.01099  ,
+                               5.10512  , 5.08007  , 5.01272  , 1.70086  , 1.69544  , 1.66979  ,
+                               3.48873  , 0.0697213, 6.67774  , 3.38621  , 0.033338 , 3.34239  ,
+                               1.79424  , 1.7281   , 5.01015  , 3.48973  , 3.42896  , 6.67795  ,
+                               3.40064  , 3.39148  , 3.34188  , 5.09069  , 1.72876  , 5.00917  ,
+                               0.119885 , 6.74841  , 3.33869  , 4.99379  , 1.69262  , 1.67183  ,
+                               0.199838 , 3.4185   , 6.67565  , 1.7213   , 5.05235  , 1.66373  ,
+                               0.21494  , 6.77616  , 6.67623  , 5.00691  , 5.05     , 1.66532  ])
+        self.assertTrue(np.allclose(coord[0, ...], coord0_std))
+
+        coord1_std = np.array([4.85582828e+00, 5.12324490e+00, 1.55763728e+00, 1.82031828e+00,
+                               1.61210490e+00, 4.91329728e+00, 5.15568828e+00, 4.91296490e+00,
+                               5.02114728e+00, 1.67640828e+00, 1.62756490e+00, 1.61183728e+00,
+                               3.41785828e+00, 6.54050490e+00, 3.42793728e+00, 3.39324828e+00,
+                               3.47558490e+00, 6.50564728e+00, 3.43286828e+00, 3.44029490e+00,
+                               3.37871728e+00, 6.60497828e+00, 3.46782490e+00, 3.42949728e+00,
+                               1.82021828e+00, 5.08114490e+00, 4.93158728e+00, 5.20431828e+00,
+                               1.80972490e+00, 5.00061728e+00, 6.56278828e+00, 6.62718490e+00,
+                               3.35101728e+00, 4.97045828e+00, 1.80536490e+00, 1.73358728e+00,
+                               6.61765828e+00, 3.43486490e+00, 6.48447728e+00, 1.57899828e+00,
+                               4.89261490e+00, 1.63632728e+00, 6.59585828e+00, 1.40657901e-01,
+                               6.51767728e+00, 3.30914005e+00, 7.86399766e-02, 6.66581642e-04])
+        self.assertTrue(np.allclose(coord[1, ...], coord1_std))
+
+if __name__ == '__main__':
+    unittest.main()
+