add to_pymatgen_structure and to_ase_structure method

Author: haidi-ustc

dpdata/deepmd/__init__.py

@@ -346,6 +346,43 @@ def from_lammps_lmp (self, file_name, type_map = None) :
             self.data = dpdata.lammps.lmp.to_system_data(lines, type_map)
         self._shift_orig_zero()
 
+    def to_pymatgen_structure(self):
+        '''
+        convert System to Pymatgen Structure obj
+
+        '''
+        structures=[]
+        try:
+           from pymatgen import Structure
+        except:
+           raise ImportError('No module pymatgen.Structure')
+
+        for system in self.to_list():
+            species=[]
+            for name,numb in zip(system.data['atom_names'],system.data['atom_numbs']):
+                species.extend([name]*numb)
+            structure=Structure(system.data['cells'][0],species,system.data['coords'][0],coords_are_cartesian=True)
+            structures.append(structure)
+        return structures
+
+    def to_ase_structure(self):
+        '''
+        convert System to ASE Atom obj
+
+        '''
+        structures=[]
+        try:
+           from ase import Atoms
+        except:
+           raise ImportError('No module ase.Atoms')
+
+        for system in self.to_list():
+            species=[]
+            for name,numb in zip(system.data['atom_names'],system.data['atom_numbs']):
+                species.extend([name]*numb)
+            structure=Atoms(symbols=species,positions=system.data['coords'][0],pbc=True,cell=system.data['cells'][0])
+            structures.append(structure)
+        return structures
 
     def to_lammps_lmp(self, file_name, frame_idx = 0) :
         """
@@ -380,6 +417,19 @@ def from_vasp_poscar(self, file_name) :
             self.data = dpdata.vasp.poscar.to_system_data(lines)
         self.rot_lower_triangular()
 
+
+    def to_vasp_string(self, frame_idx=0):
+        """
+        Dump the system in vasp POSCAR format string
+        
+        Parameters
+        ----------
+        frame_idx : int
+            The index of the frame to dump
+        """
+        assert(frame_idx < len(self.data['coords']))
+        w_str = dpdata.vasp.poscar.from_system_data(self.data, frame_idx)
+        return w_str
 
     def to_vasp_poscar(self, file_name, frame_idx = 0) :
         """
@@ -392,8 +442,7 @@ def to_vasp_poscar(self, file_name, frame_idx = 0) :
         frame_idx : int
             The index of the frame to dump
         """
-        assert(frame_idx < len(self.data['coords']))
-        w_str = dpdata.vasp.poscar.from_system_data(self.data, frame_idx)
+        w_str=self.to_vasp_string( frame_idx= frame_idx )
         with open(file_name, 'w') as fp:
             fp.write(w_str)
 

dpdata/gaussian/__init__.py

@@ -0,0 +1,29 @@
+import os
+import numpy as np
+import unittest
+from context import dpdata
+from comp_sys import CompSys
+try:
+   from ase import Atoms
+   from ase.io import write
+   exist_module=True
+except:
+   exist_module=False
+
+@unittest.skipIf(not exist_module,"skip test_ase")
+class TestASE(unittest.TestCase, CompSys):
+    
+    def setUp(self): 
+        system_1 = dpdata.System()
+        system_1.from_lammps_lmp(os.path.join('poscars', 'conf.lmp'), type_map = ['O', 'H'])
+        write('tmp.POSCAR',system_1.to_ase_structure()[0],vasp5=True)
+        self.system_1=system_1
+        self.system_2=dpdata.System('tmp.POSCAR')
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+if __name__ == '__main__':
+    unittest.main()
+    

dpdata/lammps/__init__.py

@@ -0,0 +1,28 @@
+import os
+import numpy as np
+import unittest
+from context import dpdata
+from comp_sys import CompSys
+try:
+   from pymatgen import  Structure
+   exist_module=True
+except:
+   exist_module=False
+
+@unittest.skipIf(not exist_module,"skip pymatgen")
+class TestPymatgen(unittest.TestCase, CompSys):
+    
+    def setUp(self): 
+        system_1 = dpdata.System()
+        system_1.from_lammps_lmp(os.path.join('poscars', 'conf.lmp'), type_map = ['O', 'H'])
+        system_1.to_pymatgen_structure()[0].to('poscar','tmp.POSCAR')
+        self.system_1=system_1
+        self.system_2=dpdata.System('tmp.POSCAR')
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+if __name__ == '__main__':
+    unittest.main()
+