support converting ase.Atoms to System and LabeledSystem (#290)

#190 added `from_labeled_system`, but it just passes a dict and does nothing else. This commit rewrites it.

Author: njzjz

docs/conf.py

@@ -183,4 +183,7 @@ def setup(app):
 intersphinx_mapping = {
     "numpy": ("https://docs.scipy.org/doc/numpy/", None),
     "python": ("https://docs.python.org/", None),
+    "ase": ("https://wiki.fysik.dtu.dk/ase/", None),
+    "monty": ("https://guide.materialsvirtuallab.org/monty/", None),
+    "h5py": ("https://docs.h5py.org/en/stable/", None),
 }

dpdata/plugins/ase.py

@@ -18,41 +18,101 @@ class ASEStructureFormat(Format):
     automatic detection fails.
     """
 
-    def from_labeled_system(self, data, **kwargs):
-        return data
-
-    def from_multi_systems(self, file_name, begin=None, end=None, step=None, ase_fmt=None, **kwargs):
+    def from_system(self, atoms: "ase.Atoms", **kwargs) -> dict:
+        """Convert ase.Atoms to a System.
+
+        Parameters
+        ----------
+        atoms : ase.Atoms
+            an ASE Atoms, containing a structure
+
+        Returns
+        -------
+        dict
+            data dict
+        """
+        symbols = atoms.get_chemical_symbols()
+        atom_names = list(set(symbols))
+        atom_numbs = [symbols.count(symbol) for symbol in atom_names]
+        atom_types = np.array([atom_names.index(symbol) for symbol in symbols]).astype(int)
+        cells = atoms.cell[:]
+        coords = atoms.get_positions()
+        info_dict = {
+            'atom_names': atom_names,
+            'atom_numbs': atom_numbs,
+            'atom_types': atom_types,
+            'cells': np.array([cells]).astype('float32'),
+            'coords': np.array([coords]).astype('float32'),
+            'orig': [0,0,0],
+        }
+        return info_dict
+
+    def from_labeled_system(self, atoms: "ase.Atoms", **kwargs) -> dict:
+        """Convert ase.Atoms to a LabeledSystem. Energies and forces
+        are calculated by the calculator.
+
+        Parameters
+        ----------
+        atoms : ase.Atoms
+            an ASE Atoms, containing a structure
+
+        Returns
+        -------
+        dict
+            data dict
+        
+        Raises
+        ------
+        RuntimeError
+            ASE will raise RuntimeError if the atoms does not
+            have a calculator
+        """
+        info_dict = self.from_system(atoms)
+        try:
+            energies = atoms.get_potential_energy(force_consistent=True)
+        except PropertyNotImplementedError:
+            energies = atoms.get_potential_energy()
+        forces = atoms.get_forces()
+        info_dict = {
+            ** info_dict,
+            'energies': np.array([energies]).astype('float32'),
+            'forces': np.array([forces]).astype('float32'),
+        }
+        try:
+            stress = atoms.get_stress(False)
+        except PropertyNotImplementedError:
+            pass
+        else:
+            virials = np.array([-atoms.get_volume() * stress]).astype('float32')
+            info_dict['virials'] = virials
+        return info_dict
+
+    def from_multi_systems(self, file_name: str, begin: int=None, end: int=None, step: int=None, ase_fmt: str=None, **kwargs) -> "ase.Atoms":
+        """Convert a ASE supported file to ASE Atoms.
+
+        It will finally be converted to MultiSystems.
+
+        Parameters
+        ----------
+        file_name : str
+            path to file
+        begin : int, optional
+            begin frame index
+        end : int, optional
+            end frame index
+        step : int, optional
+            frame index step
+        ase_fmt : str, optional
+            ASE format. See the ASE documentation about supported formats
+
+        Yields
+        ------
+        ase.Atoms
+            ASE atoms in the file
+        """
         frames = ase.io.read(file_name, format=ase_fmt, index=slice(begin, end, step))
         for atoms in frames:
-            symbols = atoms.get_chemical_symbols()
-            atom_names = list(set(symbols))
-            atom_numbs = [symbols.count(symbol) for symbol in atom_names]
-            atom_types = np.array([atom_names.index(symbol) for symbol in symbols]).astype(int)
-    
-            cells = atoms.cell[:]
-            coords = atoms.get_positions()
-            try:
-                energies = atoms.get_potential_energy(force_consistent=True)
-            except PropertyNotImplementedError:
-                energies = atoms.get_potential_energy()
-            forces = atoms.get_forces()
-            info_dict = {
-                'atom_names': atom_names,
-                'atom_numbs': atom_numbs,
-                'atom_types': atom_types,
-                'cells': np.array([cells]).astype('float32'),
-                'coords': np.array([coords]).astype('float32'),
-                'energies': np.array([energies]).astype('float32'),
-                'forces': np.array([forces]).astype('float32'),
-                'orig': [0,0,0],
-            }
-            try:
-                stress = atoms.get_stress(False)
-                virials = np.array([-atoms.get_volume() * stress]).astype('float32')
-                info_dict['virials'] = virials
-            except PropertyNotImplementedError:
-                pass
-            yield info_dict
+            yield atoms
 
     def to_system(self, data, **kwargs):
         '''

tests/test_to_ase.py

@@ -4,11 +4,13 @@
 from context import dpdata
 from comp_sys import CompSys, IsPBC
 try:
-   from ase import Atoms
-   from ase.io import write
-   exist_module=True
-except Exception:
-   exist_module=False
+    from ase import Atoms
+    from ase.io import write
+except ModuleNotFoundError:
+    exist_module=False
+else:
+    exist_module=True
+
 
 @unittest.skipIf(not exist_module,"skip test_ase")
 class TestASE(unittest.TestCase, CompSys, IsPBC):
@@ -24,6 +26,24 @@ def setUp(self):
         self.f_places = 6
         self.v_places = 6
 
+
+@unittest.skipIf(not exist_module, "skip test_ase")
+class TestFromASE(unittest.TestCase, CompSys, IsPBC):
+    """Test ASEStructureFormat.from_system"""
+    def setUp(self): 
+        system_1 = dpdata.System()
+        system_1.from_lammps_lmp(os.path.join('poscars', 'conf.lmp'), type_map = ['O', 'H'])
+        atoms = system_1.to_ase_structure()[0]
+        self.system_1 = system_1
+        self.system_2 = dpdata.System(atoms, fmt="ase/structure")
+        # assign the same type_map
+        self.system_2.sort_atom_names(type_map=self.system_1.get_atom_names())
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+
 if __name__ == '__main__':
     unittest.main()