diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py index 82b64e715..e7e527ff7 100644 --- a/dpdata/plugins/pymatgen.py +++ b/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 = [] diff --git a/dpdata/pymatgen/structure.py b/dpdata/pymatgen/structure.py new file mode 100644 index 000000000..b6c148da3 --- /dev/null +++ b/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 diff --git a/tests/test_from_pymatgen.py b/tests/test_from_pymatgen.py new file mode 100644 index 000000000..d3ddbe3e9 --- /dev/null +++ b/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() diff --git a/tests/test_to_pymatgen.py b/tests/test_to_pymatgen.py index d58153962..b55443d4d 100644 --- a/tests/test_to_pymatgen.py +++ b/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