diff --git a/dpdata/plugins/pymatgen.py b/dpdata/plugins/pymatgen.py index ff935d3a..b8099a3a 100644 --- a/dpdata/plugins/pymatgen.py +++ b/dpdata/plugins/pymatgen.py @@ -30,14 +30,15 @@ def to_system(self, data, **kwargs): """Convert System to Pymatgen Structure obj.""" structures = [] try: - from pymatgen.core import Structure + from pymatgen.core import Lattice, Structure except ModuleNotFoundError as e: raise ImportError("No module pymatgen.Structure") from e species = [data["atom_names"][tt] for tt in data["atom_types"]] + pbc = not (data.get("nopbc", False)) for ii in range(data["coords"].shape[0]): structure = Structure( - data["cells"][ii], + Lattice(data["cells"][ii], pbc=[pbc] * 3), species, data["coords"][ii], coords_are_cartesian=True, diff --git a/tests/test_pymatgen_structure.py b/tests/test_pymatgen_structure.py index 1b61e497..1e93829e 100644 --- a/tests/test_pymatgen_structure.py +++ b/tests/test_pymatgen_structure.py @@ -3,7 +3,7 @@ import os import unittest -from comp_sys import CompSys, IsPBC +from comp_sys import CompSys, IsNoPBC, IsPBC from context import dpdata try: @@ -42,5 +42,20 @@ def setUp(self): self.v_places = 6 +@unittest.skipIf(not exist_module, "skip pymatgen") +class TestFormToPytmatgenNopbc(unittest.TestCase, CompSys, IsNoPBC): + def setUp(self): + self.system = dpdata.System("pymatgen_data/deepmd/", fmt="deepmd/npy") + self.system.data["nopbc"] = True + self.system_1 = self.system + self.system_2 = dpdata.System().from_pymatgen_structure( + self.system.to("pymatgen/structure")[0] + ) + self.places = 6 + self.e_places = 6 + self.f_places = 6 + self.v_places = 6 + + if __name__ == "__main__": unittest.main()