Skip to content

Commit 508f5f3

Browse files
amcadmusamcadmus
authored
Merge pull request #54 from amcadmus/devel
Devel
2 parents 1fb4d32 + 3742a63 commit 508f5f3

File tree

4 files changed

+116
-7
lines changed

4 files changed

+116
-7
lines changed

‎dpdata/deepmd/comp.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ def dump(folder,
6666
raise RuntimeError('found ' + str(sets) + ' in ' + folder + 'not a clean deepmd raw dir. please firstly clean set.* then try compress')
6767
# dump raw
6868
np.savetxt(os.path.join(folder, 'type.raw'), data['atom_types'], fmt = '%d')
69+
np.savetxt(os.path.join(folder, 'type_map.raw'), data['atom_names'], fmt = '%s')
6970
# reshape frame properties and convert prec
7071
nframes = data['cells'].shape[0]
7172
cells = np.reshape(data['cells'], [nframes, 9]).astype(comp_prec)

‎dpdata/deepmd/raw.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,22 @@ def load_type(folder, type_map = None) :
1010
for ii in range (ntypes) :
1111
data['atom_numbs'].append(np.count_nonzero(data['atom_types'] == ii))
1212
data['atom_names'] = []
13-
if type_map == None :
13+
# if find type_map.raw, use it
14+
if os.path.isfile(os.path.join(folder, 'type_map.raw')) :
15+
with open(os.path.join(folder, 'type_map.raw')) as fp:
16+
my_type_map = fp.read().split()
17+
# else try to use arg type_map
18+
elif type_map is not None:
19+
my_type_map = type_map
20+
# in the last case, make artificial atom names
21+
else:
22+
my_type_map = []
1423
for ii in range(ntypes) :
15-
data['atom_names'].append('Type_%d' % ii)
16-
else :
17-
assert(len(type_map) >= len(data['atom_numbs']))
18-
for ii in range(len(data['atom_numbs'])) :
19-
data['atom_names'].append(type_map[ii])
24+
my_type_map.append('Type_%d' % ii)
25+
assert(len(my_type_map) >= len(data['atom_numbs']))
26+
for ii in range(len(data['atom_numbs'])) :
27+
data['atom_names'].append(my_type_map[ii])
28+
2029
return data
2130

2231

@@ -49,6 +58,7 @@ def dump (folder, data) :
4958
os.makedirs(folder, exist_ok = True)
5059
nframes = data['cells'].shape[0]
5160
np.savetxt(os.path.join(folder, 'type.raw'), data['atom_types'], fmt = '%d')
61+
np.savetxt(os.path.join(folder, 'type_map.raw'), data['atom_names'], fmt = '%s')
5262
np.savetxt(os.path.join(folder, 'box.raw'), np.reshape(data['cells'], [nframes, 9]))
5363
np.savetxt(os.path.join(folder, 'coord.raw'), np.reshape(data['coords'], [nframes, -1]))
5464
if 'energies' in data :

‎setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
readme = f.read()
1313

1414
# install_requires = ['xml']
15-
install_requires=['numpy>=1.14.3', 'monty<2.0.7']
15+
install_requires=['numpy>=1.14.3', 'monty']
1616

1717
setuptools.setup(
1818
name="dpdata",

‎tests/test_deepmd_raw.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,104 @@ def tearDown(self) :
3333
shutil.rmtree('tmp.deepmd')
3434

3535

36+
class TestDeepmdTypeMap(unittest.TestCase):
37+
def tearDown(self) :
38+
if os.path.exists('tmp.deepmd'):
39+
shutil.rmtree('tmp.deepmd')
40+
41+
def test_type_map (self) :
42+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
43+
fmt = 'vasp/outcar')
44+
system_1.to_deepmd_raw('tmp.deepmd')
45+
with open(os.path.join('tmp.deepmd', 'type_map.raw')) as fp:
46+
tm = fp.read().split()
47+
self.assertEqual(tm, ['O', 'H'])
48+
self.assertEqual(system_1['atom_names'], ['O', 'H'])
49+
self.assertEqual(system_1['atom_types'][0], 0)
50+
self.assertEqual(system_1['atom_types'][1], 0)
51+
self.assertEqual(system_1['atom_types'][2], 1)
52+
self.assertEqual(system_1['atom_types'][3], 1)
53+
self.assertEqual(system_1['atom_types'][4], 1)
54+
self.assertEqual(system_1['atom_types'][5], 1)
55+
56+
def test_type_map_load (self) :
57+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
58+
fmt = 'vasp/outcar')
59+
system_1.to_deepmd_raw('tmp.deepmd')
60+
system_2 = dpdata.LabeledSystem('tmp.deepmd')
61+
self.assertEqual(system_2['atom_names'], ['O', 'H'])
62+
self.assertEqual(system_2['atom_types'][0], 0)
63+
self.assertEqual(system_2['atom_types'][1], 0)
64+
self.assertEqual(system_2['atom_types'][2], 1)
65+
self.assertEqual(system_2['atom_types'][3], 1)
66+
self.assertEqual(system_2['atom_types'][4], 1)
67+
self.assertEqual(system_2['atom_types'][5], 1)
68+
self.assertEqual(system_2['atom_numbs'][0], 2)
69+
self.assertEqual(system_2['atom_numbs'][1], 4)
70+
71+
def test_type_map_enforce (self) :
72+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
73+
fmt = 'vasp/outcar')
74+
system_1.to_deepmd_raw('tmp.deepmd')
75+
system_2 = dpdata.LabeledSystem('tmp.deepmd', type_map = ['H', 'O'])
76+
self.assertEqual(system_2['atom_names'], ['H', 'O'])
77+
self.assertEqual(system_2['atom_types'][0], 1)
78+
self.assertEqual(system_2['atom_types'][1], 1)
79+
self.assertEqual(system_2['atom_types'][2], 0)
80+
self.assertEqual(system_2['atom_types'][3], 0)
81+
self.assertEqual(system_2['atom_types'][4], 0)
82+
self.assertEqual(system_2['atom_types'][5], 0)
83+
self.assertEqual(system_2['atom_numbs'][0], 4)
84+
self.assertEqual(system_2['atom_numbs'][1], 2)
85+
86+
def test_npy_type_map (self) :
87+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
88+
fmt = 'vasp/outcar')
89+
system_1.to_deepmd_npy('tmp.deepmd')
90+
with open(os.path.join('tmp.deepmd', 'type_map.raw')) as fp:
91+
tm = fp.read().split()
92+
self.assertEqual(tm, ['O', 'H'])
93+
self.assertEqual(system_1['atom_names'], ['O', 'H'])
94+
self.assertEqual(system_1['atom_types'][0], 0)
95+
self.assertEqual(system_1['atom_types'][1], 0)
96+
self.assertEqual(system_1['atom_types'][2], 1)
97+
self.assertEqual(system_1['atom_types'][3], 1)
98+
self.assertEqual(system_1['atom_types'][4], 1)
99+
self.assertEqual(system_1['atom_types'][5], 1)
100+
101+
def test_npy_type_map_load (self) :
102+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
103+
fmt = 'vasp/outcar')
104+
system_1.to_deepmd_npy('tmp.deepmd')
105+
system_2 = dpdata.LabeledSystem('tmp.deepmd', fmt = 'deepmd/npy')
106+
self.assertEqual(system_2['atom_names'], ['O', 'H'])
107+
self.assertEqual(system_2['atom_types'][0], 0)
108+
self.assertEqual(system_2['atom_types'][1], 0)
109+
self.assertEqual(system_2['atom_types'][2], 1)
110+
self.assertEqual(system_2['atom_types'][3], 1)
111+
self.assertEqual(system_2['atom_types'][4], 1)
112+
self.assertEqual(system_2['atom_types'][5], 1)
113+
self.assertEqual(system_2['atom_numbs'][0], 2)
114+
self.assertEqual(system_2['atom_numbs'][1], 4)
115+
116+
def test_npy_type_map_enforce (self) :
117+
system_1 = dpdata.LabeledSystem('poscars/OUTCAR.h2o.md',
118+
fmt = 'vasp/outcar')
119+
system_1.to_deepmd_npy('tmp.deepmd')
120+
system_2 = dpdata.LabeledSystem('tmp.deepmd', type_map = ['H', 'O'], fmt = 'deepmd/npy')
121+
self.assertEqual(system_2['atom_names'], ['H', 'O'])
122+
self.assertEqual(system_2['atom_types'][0], 1)
123+
self.assertEqual(system_2['atom_types'][1], 1)
124+
self.assertEqual(system_2['atom_types'][2], 0)
125+
self.assertEqual(system_2['atom_types'][3], 0)
126+
self.assertEqual(system_2['atom_types'][4], 0)
127+
self.assertEqual(system_2['atom_types'][5], 0)
128+
self.assertEqual(system_2['atom_numbs'][0], 4)
129+
self.assertEqual(system_2['atom_numbs'][1], 2)
130+
131+
132+
133+
36134
class TestDeepmdRawNoLabels(unittest.TestCase, CompSys) :
37135
def setUp (self) :
38136
self.system_1 = dpdata.System('poscars/POSCAR.h2o.md',

0 commit comments

Comments
 (0)