diff --git a/dpdata/lammps/dump.py b/dpdata/lammps/dump.py index fe549b95..89e75e4d 100644 --- a/dpdata/lammps/dump.py +++ b/dpdata/lammps/dump.py @@ -355,6 +355,62 @@ def split_traj(dump_lines): return None +def from_system_data(system, f_idx=0, timestep=0): + """Convert system data to LAMMPS dump format string. + + Parameters + ---------- + system : dict + System data dictionary containing atoms, coordinates, cell, etc. + f_idx : int, optional + Frame index to dump (default: 0) + timestep : int, optional + Timestep number for the dump (default: 0) + + Returns + ------- + str + LAMMPS dump format string + """ + ret = "" + + # Get basic system info + natoms = sum(system["atom_numbs"]) + coords = system["coords"][f_idx] + cell = system["cells"][f_idx] + atom_types = system["atom_types"] + orig = system.get("orig", np.zeros(3)) + + # Convert cell to dump format (bounds and tilt) + bounds, tilt = box2dumpbox(orig, cell) + + # Write timestep + ret += "ITEM: TIMESTEP\n" + ret += f"{timestep}\n" + + # Write number of atoms + ret += "ITEM: NUMBER OF ATOMS\n" + ret += f"{natoms}\n" + + # Write box bounds + ret += "ITEM: BOX BOUNDS xy xz yz pp pp pp\n" + ret += f"{bounds[0][0]:.10f} {bounds[0][1]:.10f} {tilt[0]:.10f}\n" + ret += f"{bounds[1][0]:.10f} {bounds[1][1]:.10f} {tilt[1]:.10f}\n" + ret += f"{bounds[2][0]:.10f} {bounds[2][1]:.10f} {tilt[2]:.10f}\n" + + # Write atoms header + ret += "ITEM: ATOMS id type x y z\n" + + # Write atom data + for ii in range(natoms): + atom_id = ii + 1 # LAMMPS uses 1-based indexing + atom_type = atom_types[ii] + 1 # LAMMPS uses 1-based type indexing + x, y, z = coords[ii] + ret += f"{atom_id} {atom_type} {x:.10f} {y:.10f} {z:.10f}\n" + + return ret + + if __name__ == "__main__": # fname = 'dump.hti' # lines = open(fname).read().split('\n') diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py index 9949e99e..b00d4ff0 100644 --- a/dpdata/plugins/lammps.py +++ b/dpdata/plugins/lammps.py @@ -170,3 +170,24 @@ def from_system( ) register_spin(data) return data + + def to_system(self, data, file_name: FileType, frame_idx=0, timestep=0, **kwargs): + """Dump the system in LAMMPS dump format. + + Parameters + ---------- + data : dict + System data + file_name : str + The output file name + frame_idx : int + The index of the frame to dump + timestep : int + The timestep number for the dump + **kwargs : dict + other parameters + """ + assert frame_idx < len(data["coords"]) + w_str = dpdata.lammps.dump.from_system_data(data, frame_idx, timestep) + with open_file(file_name, "w") as fp: + fp.write(w_str) diff --git a/tests/test_lammps_lmp_dump.py b/tests/test_lammps_lmp_dump.py index 4d593211..c2c2f811 100644 --- a/tests/test_lammps_lmp_dump.py +++ b/tests/test_lammps_lmp_dump.py @@ -30,6 +30,16 @@ def setUp(self): self.system.from_fmt("tmp.lmp", fmt="lammps/lmp", type_map=["O", "H"]) +class TestDumpToFunc(unittest.TestCase, TestPOSCARoh): + def setUp(self): + tmp_system = dpdata.System( + os.path.join("poscars", "conf.lmp"), type_map=["O", "H"] + ) + tmp_system.to("lammps/dump", "tmp.dump") + self.system = dpdata.System() + self.system.from_fmt("tmp.dump", fmt="lammps/dump", type_map=["O", "H"]) + + class TestLmpRotateTriAngle(unittest.TestCase): def test_simple_cubic(self): cubic_cell = np.diag([5, 5, 5])