Skip to content

fix: serialize LAMMPS HDF5 trajectories #381

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
njzjz-bot wants to merge 1 commit into
base: master
Choose a base branch
from
+35 −1
Open

fix: serialize LAMMPS HDF5 trajectories #381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion dpgen2/op/run_lmp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ def execute(

ret_dict = {
"log": work_dir / lmp_log_name,
"traj": work_dir / lmp_traj_name,
"traj": self.get_traj(work_dir / lmp_traj_name),
"model_devi": self.get_model_devi(work_dir / lmp_model_devi_name),
}
plm_output = (
Expand All @@ -223,6 +223,10 @@ def execute(
def get_model_devi(self, model_devi_file):
return model_devi_file

def get_traj(self, traj_file):
"""Return the filesystem trajectory used by the standard backend."""
return traj_file

@staticmethod
def lmp_args():
doc_lmp_cmd = "The command of LAMMPS"
Expand Down Expand Up @@ -412,3 +416,7 @@ def get_output_sign(cls):

def get_model_devi(self, model_devi_file):
return np.loadtxt(model_devi_file)

@wanghan-iapcm wanghan-iapcm Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has the identical defect and needs to move in the same change, not a follow-up. np.loadtxt returns an ndarray, which is not a list or dict, so flatten drops it too — keys = [].

The reason it cannot wait: if only traj becomes a list, the two artifacts arrive downstream with different lengths, and dpgen2/exploration/selector/conf_selector_frame.py:88-89 is

ntraj = len(trajs)
assert ntraj == len(model_devis)

I simulated all three states through the real handle_output_artifact + handle_input_artifact round trip with two tasks:

both scalar (current head)      trajs=0  model_devis=0   assert passes vacuously, silent data loss
traj list, model_devi scalar    trajs=2  model_devis=0   AssertionError mid-workflow
both lists                      trajs=2  model_devis=2   correct

So fixing traj alone is worse than fixing neither. return [np.loadtxt(model_devi_file)] alongside the traj change.


def get_traj(self, traj_file):
"""Return trajectory text for serialization into an HDF5 dataset."""
return traj_file.read_text()

@wanghan-iapcm wanghan-iapcm Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the line to change. read_text() returns a str, which is exactly as much a scalar as the Path it replaced, so flatten drops it and the dataset loop never runs. Confirmed against the installed dflow: flatten('trajectory data') is {}, and handle_output_artifact with Artifact(HDF5Datasets) and an int slices produces an .h5 with keys = [].

Suggest returning a list containing the Path, not the text:

def get_traj(self, traj_file):
    return [traj_file]

The list is what flatten needs. Keeping it a Path also routes through dflow's own Path branch, which read_text() bypasses:

if v.is_file():
    try:
        data = v.read_text(encoding="utf-8"); dtype = "utf-8"
    except Exception:
        data = np.void(v.read_bytes()); dtype = "binary"
    d = f.create_dataset(s, data=data)
    d.attrs["type"] = "file"; d.attrs["path"] = str(v); d.attrs["dtype"] = dtype

So you get the is_file() guard, a binary fallback for a non-UTF-8 dump, and the type/path attrs for free. As written, a missing dump raises a bare FileNotFoundError out of execute() and a non-UTF-8 dump raises UnicodeDecodeError — I reproduced both. Those are unlikely in practice since dpgen2 generates the dump directive itself, but there is no reason to give up the guards. [Path] still arrives at the consumer as decoded text, because HDF5Dataset.get_data() decodes on dtype == "utf-8" in both branches.

The docstring on this method also needs updating: "Return trajectory text for serialization into an HDF5 dataset" asserts that returning the text is what causes serialization, and that is what is not true. Note too that RunLmpHDF5 does not override execute, so it inherits RunLmp.execute's Returns section, which still documents traj and model_devi as Artifact(Path).

26 changes: 26 additions & 0 deletions tests/op/test_run_lmp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
)
from dpgen2.op.run_lmp import (
RunLmp,
RunLmpHDF5,
get_ele_temp,
merge_pimd_files,
set_models,
Expand Down Expand Up @@ -149,6 +150,31 @@ def test_extra_outputs(self):
"Hello -i in.lammps -log log.lammps",
)

@patch("dpgen2.op.run_lmp.run_command")
def test_hdf5_outputs_dataset_values(self, mocked_run):
"""Return serializable data instead of filesystem paths for HDF5."""

def write_outputs(*args, **kwargs):
Path(lmp_traj_name).write_text("trajectory data")
np.savetxt(lmp_model_devi_name, np.arange(7).reshape(1, 7))
return 0, "foo\n", ""

mocked_run.side_effect = write_outputs

out = RunLmpHDF5().execute(
OPIO(
{
"config": {"command": "mylmp"},
"task_name": self.task_name,
"task_path": self.task_path,
"models": self.models,
}
)
)

self.assertEqual(out["traj"], "trajectory data")

@wanghan-iapcm wanghan-iapcm Sep 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test is not vacuous — I checked, reverting get_traj to the bare Path makes it fail with AssertionError: PosixPath('task_000/traj.dump') != 'trajectory data'. But it asserts one layer above the defect, so it is green while the artifact it exists to protect is empty.

out["traj"] is the OP's in-memory return value. The failure lives in dflow's serialization of that value, which this test never invokes. I fed the exact value the test asserts on into the real handle_output_artifact and got an .h5 with zero datasets.

Issue #355 asked for "a test for RunLmpHDF5.execute output types". The letter is satisfied; the contract is not. A round-trip assertion is what would have caught this and what would stop it regressing:

from dflow.python.utils import handle_output_artifact
handle_output_artifact("traj", out["traj"], Artifact(HDF5Datasets), slices=0, data_root=tmp)
# then open the produced .h5 and assert its key set is non-empty

Worth doing for model_devi in the same test. Minor, while you are here: this lives in TestRunLmp but exercises RunLmpHDF5; a separate TestRunLmpHDF5 class would read better.

np.testing.assert_array_equal(out["model_devi"], np.arange(7))


class TestRunLmpDist(unittest.TestCase):
lmp_config = """variable NSTEPS equal 1000
Expand Down