PyMatgenStructureFormat with from_system method

dpdata/plugins/pymatgen.py

@@ -1,11 +1,29 @@
 import numpy as np
 
 import dpdata.pymatgen.molecule
+import dpdata.pymatgen.structure
 from dpdata.format import Format
 
 
 @Format.register("pymatgen/structure")
 class PyMatgenStructureFormat(Format):
+    def from_system(self, structure, **kwargs) -> dict:
+        """Convert pymatgen.core.Structure to System.
+
+        Parameters
+        ----------
+        structure : pymatgen.core.Structure
+            a Pymatgen Structure, containing a structure
+        **kwargs : dict
+            other parameters
+
+        Returns
+        -------
+        dict
+            data dict
+        """
+        return dpdata.pymatgen.structure.from_system_data(structure)
+
     def to_system(self, data, **kwargs):
         """Convert System to Pymatgen Structure obj."""
         structures = []

dpdata/pymatgen/structure.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+try:
+    from pymatgen.core import Structure  # noqa: F401
+except ImportError:
+    pass
+
+
+def from_system_data(structure) -> dict:
+    symbols = [site.species_string for site in structure]
+    atom_names = list(structure.symbol_set)
+    atom_numbs = [symbols.count(symbol) for symbol in atom_names]
+    atom_types = np.array([atom_names.index(symbol) for symbol in symbols]).astype(int)
+    coords = structure.cart_coords
+    cells = structure.lattice.matrix
+
+    info_dict = {
+        "atom_names": atom_names,
+        "atom_numbs": atom_numbs,
+        "atom_types": atom_types,
+        "coords": np.array([coords]),
+        "cells": np.array([cells]),
+    }
+    return info_dict

tests/test_from_pymatgen.py

@@ -0,0 +1,30 @@
+import os
+import unittest
+
+from comp_sys import CompSys
+from context import dpdata
+
+try:
+    from pymatgen.core import Structure  # noqa: F401
+
+    exist_module = True
+except Exception:
+    exist_module = False
+
+
+@unittest.skipIf(not exist_module, "skip pymatgen")
+class TestFormPytmatgen(unittest.TestCase, CompSys):
+    def setUp(self):
+        structure = Structure.from_file(os.path.join("poscars", "POSCAR.P42nmc"))
+        self.system_1 = dpdata.System(structure, fmt="pymatgen/structure")
+        self.system_2 = dpdata.System(
+            os.path.join("poscars", "POSCAR.P42nmc"), fmt="poscar"
+        )
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 6
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_to_pymatgen.py

@@ -5,7 +5,7 @@
 from context import dpdata
 
 try:
-    from pymatgen import Structure  # noqa: F401
+    from pymatgen.core import Structure  # noqa: F401
 
     exist_module = True
 except Exception:
@@ -19,7 +19,7 @@ def setUp(self):
         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")
+        system_1.to_pymatgen_structure()[0].to(filename="tmp.POSCAR", fmt="poscar")
         self.system_1 = system_1
         self.system_2 = dpdata.System("tmp.POSCAR")
         self.places = 6