diff --git a/dpdata/lammps/lmp.py b/dpdata/lammps/lmp.py index e259aa5c..33cddb41 100644 --- a/dpdata/lammps/lmp.py +++ b/dpdata/lammps/lmp.py @@ -484,7 +484,7 @@ def rotate_to_lower_triangle( return cell, coord -def from_system_data(system, f_idx=0): +def from_system_data(system, f_idx=0, output_spins=False): ret = "" ret += "\n" natoms = sum(system["atom_numbs"]) @@ -529,7 +529,7 @@ def from_system_data(system, f_idx=0): + "\n" ) # noqa: UP031 - if "spins" in system: + if "spins" in system and output_spins: coord_fmt = ( coord_fmt.strip("\n") + " " @@ -544,7 +544,7 @@ def from_system_data(system, f_idx=0): ) # noqa: UP031 spins_norm = np.linalg.norm(system["spins"][f_idx], axis=1) for ii in range(natoms): - if "spins" in system: + if "spins" in system and output_spins: if spins_norm[ii] != 0: ret += coord_fmt % ( ii + 1, diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py index b00d4ff0..0db58b80 100644 --- a/dpdata/plugins/lammps.py +++ b/dpdata/plugins/lammps.py @@ -123,7 +123,10 @@ def to_system(self, data, file_name: FileType, frame_idx=0, **kwargs): other parameters """ assert frame_idx < len(data["coords"]) - w_str = dpdata.lammps.lmp.from_system_data(data, frame_idx) + output_spins = bool(kwargs.pop("output_spins", False)) + w_str = dpdata.lammps.lmp.from_system_data( + data, frame_idx, output_spins=output_spins + ) with open_file(file_name, "w") as fp: fp.write(w_str) diff --git a/tests/test_lammps_spin.py b/tests/test_lammps_spin.py index d3d58920..bf2471ec 100644 --- a/tests/test_lammps_spin.py +++ b/tests/test_lammps_spin.py @@ -48,8 +48,8 @@ def setUp(self): def tearDown(self): pass # if os.path.isfile(self.lmp_coord_name):os.remove(self.lmp_coord_name) - def test_dump_input(self): - self.tmp_system.to("lammps/lmp", self.lmp_coord_name) + def test_dump_input(self, output_spins=False): + self.tmp_system.to("lammps/lmp", self.lmp_coord_name, output_spins=output_spins) self.assertTrue(os.path.isfile(self.lmp_coord_name)) with open(self.lmp_coord_name) as f: c = f.read() @@ -58,9 +58,9 @@ def test_dump_input(self): 2 2 1.2621856000 0.7018028000 0.5513885000 0.0000000000 0.8000000000 0.6000000000 5.0000000000""" self.assertTrue(coord_ref in c) - def test_dump_input_zero_spin(self): + def test_dump_input_zero_spin(self, output_spins=True): self.tmp_system.data["spins"] = [[[0, 0, 0], [0, 0, 0]]] - self.tmp_system.to("lammps/lmp", self.lmp_coord_name) + self.tmp_system.to("lammps/lmp", self.lmp_coord_name, output_spins=output_spins) self.assertTrue(os.path.isfile(self.lmp_coord_name)) with open(self.lmp_coord_name) as f: c = f.read()