From 58229ab321244abacd54a7038de65ba1de29ba7f Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Tue, 14 Jul 2026 16:32:23 +0800 Subject: [PATCH 1/8] use template as submit.sh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 增加了 scf-setup 中复制提交脚本的功能（原本只能让pesmaker自动生成）；如果启用了脚本复制功能，submit.sh 会按照给定的模板直接复制，如果没有启用原样复制功能，则按照 pesmaker 原本的设置生成 submit.sh 脚本 --- docs/commands/scf-setup.md | 22 ++++++ docs/commands/submit.md | 5 ++ src/pesmaker/jobs/scripts.py | 20 +++++ tests/test_job_scripts.py | 148 +++++++++++++++++++++++++++++++++++ 4 files changed, 195 insertions(+) create mode 100644 tests/test_job_scripts.py diff --git a/docs/commands/scf-setup.md b/docs/commands/scf-setup.md index c35c5a4..2026629 100644 --- a/docs/commands/scf-setup.md +++ b/docs/commands/scf-setup.md @@ -29,6 +29,28 @@ jobs: sub_file: templates/sbatch/vasp_cpu_36.sh ``` +## Copy A Submit Template Verbatim + +By default, `sub_file` is a renderable template: PESMaker replaces supported +placeholders and refreshes the Slurm job name. To copy the template directly +to every SCF folder as `submit.sh`, enable: + +```yaml +jobs: + sub_file: templates/sbatch/vasp_cpu_36.sh + copy_sub_file: true +``` + +In this mode PESMaker performs a byte-for-byte copy. It does not replace +`{command}`, `{job_name}`, resource placeholders, literal Slurm job names, or +literal command lines. The template must therefore be ready to submit without +per-job rendering. `copy_sub_file` must be `true` or `false` and defaults to +`false`; it applies only to VASP SCF/labeling submit scripts. + +If `pesmaker submit` refreshes a pending or retry SCF script, copy mode copies +the current template again instead of rendering it. Edit the source template, +not an already prepared `submit.sh`, when this mode is enabled. + ## Inputs PESMaker looks for structures in this order: diff --git a/docs/commands/submit.md b/docs/commands/submit.md index 41aeaec..7a1e66f 100644 --- a/docs/commands/submit.md +++ b/docs/commands/submit.md @@ -113,6 +113,11 @@ happens during `--dry-run`, while the scheduler itself is not called. If no script-refresh setting is present, PESMaker submits the existing `submit.sh` without rewriting it. +For SCF jobs, `jobs.copy_sub_file: true` changes that refresh into a direct, +byte-for-byte copy of the configured template. No placeholders, Slurm job +names, resource values, or commands are rewritten. The option requires a +submit template and defaults to `false`. + When `jobs.sub_file` is provided, PESMaker does not scan and rewrite literal resource directives such as `#SBATCH --ntasks` or literal VASP command lines. It only replaces explicit placeholders such as `{command}`, `{job_name}`, diff --git a/src/pesmaker/jobs/scripts.py b/src/pesmaker/jobs/scripts.py index df62435..ad2781c 100644 --- a/src/pesmaker/jobs/scripts.py +++ b/src/pesmaker/jobs/scripts.py @@ -19,6 +19,7 @@ from __future__ import annotations import re +import shutil from pathlib import Path from typing import Any @@ -45,6 +46,17 @@ def _write_submit_script( resources: JobResources | None = None, ) -> Path: template_path = _job_template_path(config, stage) + if _copy_sub_file_enabled(config, stage): + if template_path is None: + raise ValueError( + "jobs.copy_sub_file requires an SCF submit template in " + "jobs.sub_file, jobs.sbatch_templates, or jobs.sbatch_template" + ) + path = workdir / "submit.sh" + if template_path.resolve() != path.resolve(): + shutil.copy2(template_path, path) + return path + job_name = _submit_job_name(workdir) resources = resources or _job_resources(config) engine = _stage_engine(config, stage) @@ -97,6 +109,14 @@ def _write_submit_script( return path +def _copy_sub_file_enabled(config: PESMakerConfig, stage: str) -> bool: + """Return whether an SCF submit template should be copied verbatim.""" + value = config.jobs.options.get("copy_sub_file", False) + if not isinstance(value, bool): + raise ValueError("jobs.copy_sub_file must be true or false") + return stage == "labeling" and value + + def _job_template_path(config: PESMakerConfig, stage: str) -> Path | None: sub_file = config.jobs.options.get("sub_file") if isinstance(sub_file, dict): diff --git a/tests/test_job_scripts.py b/tests/test_job_scripts.py new file mode 100644 index 0000000..bd5c9b5 --- /dev/null +++ b/tests/test_job_scripts.py @@ -0,0 +1,148 @@ +# Copyright 2026 Ting Liang and PESMaker development team +# This file is part of PESMaker. +# +# PESMaker is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. + +"""Tests for submission script rendering and literal template copies.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from pesmaker.config.schema import PESMakerConfig +from pesmaker.jobs.scripts import _write_submit_script +from pesmaker.jobs.submit import submit_jobs + + +def _config(template: Path | None, *, copy_sub_file: object = False): + jobs: dict[str, object] = { + "copy_sub_file": copy_sub_file, + "cores_cpu": 36, + } + if template is not None: + jobs["sub_file"] = str(template) + return PESMakerConfig.from_mapping( + { + "project": "submit_script_test", + "labeling": {"engine": "vasp"}, + "jobs": jobs, + } + ) + + +def test_scf_submit_template_can_be_copied_verbatim(tmp_path): + """The opt-in copy mode must not render or normalize template bytes.""" + template = tmp_path / "template.sh" + template_bytes = ( + b"#!/bin/bash\r\n" + b"#SBATCH --job-name=keep-this-name\r\n" + b"#SBATCH --ntasks=7\r\n" + b"{command}\r\n" + ) + template.write_bytes(template_bytes) + workdir = tmp_path / "labeling" / "selected_000001" + workdir.mkdir(parents=True) + + path = _write_submit_script( + _config(template, copy_sub_file=True), + workdir, + stage="labeling", + command="/new/path/vasp_std", + ) + + assert path == workdir / "submit.sh" + assert path.read_bytes() == template_bytes + + +def test_scf_submit_template_uses_existing_rendering_by_default(tmp_path): + """Omitting copy_sub_file must preserve the original rendering behavior.""" + template = tmp_path / "template.sh" + template.write_text( + "#!/bin/bash\n#SBATCH --job-name=old-name\n{command}\n", + encoding="utf-8", + ) + workdir = tmp_path / "labeling" / "selected_000001" + workdir.mkdir(parents=True) + + path = _write_submit_script( + _config(template), + workdir, + stage="labeling", + command="/new/path/vasp_std", + ) + text = path.read_text(encoding="utf-8") + + assert "#SBATCH --job-name=labeling_sel000001" in text + assert "mpirun -np 36 /new/path/vasp_std" in text + + +def test_scf_submit_copy_requires_a_template(tmp_path): + """Copy mode should fail clearly instead of falling back to generated text.""" + workdir = tmp_path / "labeling" / "selected_000001" + workdir.mkdir(parents=True) + + with pytest.raises(ValueError, match="copy_sub_file requires"): + _write_submit_script( + _config(None, copy_sub_file=True), + workdir, + stage="labeling", + command="vasp_std", + ) + + +def test_scf_submit_copy_switch_must_be_boolean(tmp_path): + """String values should not silently enable or disable copy mode.""" + template = tmp_path / "template.sh" + template.write_text("#!/bin/bash\n", encoding="utf-8") + workdir = tmp_path / "labeling" / "selected_000001" + workdir.mkdir(parents=True) + + with pytest.raises(ValueError, match="must be true or false"): + _write_submit_script( + _config(template, copy_sub_file="true"), + workdir, + stage="labeling", + command="vasp_std", + ) + + +def test_scf_submit_refresh_respects_verbatim_copy_mode(tmp_path): + """Submitting a pending SCF job must not render a copied template later.""" + template = tmp_path / "template.sh" + template_bytes = b"#!/bin/bash\r\n#SBATCH -J unchanged\r\n{command}\r\n" + template.write_bytes(template_bytes) + output_dir = tmp_path / "labeling" + workdir = output_dir / "selected_000001" + workdir.mkdir(parents=True) + (workdir / "submit.sh").write_text("old script\n", encoding="utf-8") + (output_dir / "labeling_manifest.jsonl").write_text( + json.dumps({"workdir": str(workdir)}) + "\n", + encoding="utf-8", + ) + config = PESMakerConfig.from_mapping( + { + "project": "submit_refresh_test", + "labeling": { + "engine": "vasp", + "output_dir": str(output_dir), + "command": "/new/path/vasp_std", + }, + "jobs": { + "sub_file": str(template), + "copy_sub_file": True, + "cores_cpu": 36, + }, + } + ) + + submit_jobs(config, dry_run=True) + + assert (workdir / "submit.sh").read_bytes() == template_bytes + log = (output_dir / "scf_submitted_jobs.txt").read_text(encoding="utf-8") + assert "REFRESHED submit script" in log From b0562144c115948cec250273c0f4e0efc57d96f6 Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Wed, 15 Jul 2026 14:41:43 +0800 Subject: [PATCH 2/8] updated background submission --- docs/commands/submit.md | 39 ++++++++++ src/pesmaker/cli.py | 23 +++++- src/pesmaker/jobs/submit.py | 65 ++++++++++++++++ tests/test_submit_background.py | 133 ++++++++++++++++++++++++++++++++ 4 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 tests/test_submit_background.py diff --git a/docs/commands/submit.md b/docs/commands/submit.md index 7a1e66f..9933913 100644 --- a/docs/commands/submit.md +++ b/docs/commands/submit.md @@ -58,6 +58,45 @@ incomplete VASP output, or an electronic SCF nonconvergence marker. If `Jobs to submit` is `0`, all discovered calculations are complete according to the configured checks. +## Run the Whole Submission in the Background + +Use `--background` to detach the whole PESMaker submit operation from the +current terminal or SSH session: + +```bash +pesmaker submit run.yaml --background +``` + +PESMaker prints the detached process ID and a unique log path, then returns to +the shell immediately: + +```text +Background submission started. +Process ID : 18420 +Log : /path/to/labeling/scf_submit_20260715-153000_18419.log +``` + +This option protects the outer PESMaker process. It is especially useful with +`jobs.submit_command: bash`, because the detached worker keeps the complete +serial sequence alive and starts the next calculation only after the current +`submit.sh` exits. Closing the terminal does not stop that worker. + +Monitor it with the printed log path or the normal calculation output files: + +```bash +tail -f labeling/scf_submit_20260715-153000_18419.log +``` + +The same option works for other stages: + +```bash +pesmaker submit run.yaml --stage sampling --background +pesmaker submit run.yaml --stage training --background +``` + +`--background` cannot be combined with `--dry-run`. Preview in the foreground +first, then start the real background submission. + ## Stage Names `submit` is one command with different stages: diff --git a/src/pesmaker/cli.py b/src/pesmaker/cli.py index e0e5431..a4f2bc0 100644 --- a/src/pesmaker/cli.py +++ b/src/pesmaker/cli.py @@ -38,7 +38,7 @@ format_generate_summary, generate_structures, ) -from pesmaker.jobs.submit import submit_jobs +from pesmaker.jobs.submit import start_background_submit, submit_jobs from pesmaker.labelers.vasp import setup_labeling from pesmaker.plot import available_plot_commands, run_plot_command from pesmaker.results import StageResult @@ -198,6 +198,14 @@ def main(argv: list[str] | None = None) -> int: action="store_true", help="Write the submission commands without calling the scheduler.", ) + submit_parser.add_argument( + "--background", + action="store_true", + help=( + "Run the whole submit operation in a detached process that " + "survives terminal or SSH disconnection." + ), + ) init_parser = subparsers.add_parser( "init", @@ -288,6 +296,19 @@ def main(argv: list[str] | None = None) -> int: return 0 if args.command == "submit": + if args.background: + if args.dry_run: + raise ValueError("--background cannot be used with --dry-run") + process = start_background_submit( + config, + args.config, + stage=args.stage, + ) + print("Background submission started.") + print(f"Process ID : {process.pid}") + print(f"Log : {process.log_path}") + print() + return 0 _print_submit_result( submit_jobs(config, stage=args.stage, dry_run=args.dry_run), config_path=args.config, diff --git a/src/pesmaker/jobs/submit.py b/src/pesmaker/jobs/submit.py index da7eeff..af04ed7 100644 --- a/src/pesmaker/jobs/submit.py +++ b/src/pesmaker/jobs/submit.py @@ -19,8 +19,11 @@ from __future__ import annotations from dataclasses import dataclass +from datetime import datetime +import os import shlex import subprocess +import sys from pathlib import Path from pesmaker.artifacts import _read_manifest, _section_output_dir @@ -55,6 +58,68 @@ class PreparedJob: submit_script: Path | None = None +@dataclass(frozen=True) +class BackgroundSubmitProcess: + """Detached PESMaker submit process and its combined output log.""" + + pid: int + log_path: Path + + +def start_background_submit( + config: PESMakerConfig, + config_path: Path, + *, + stage: str = "scf", +) -> BackgroundSubmitProcess: + """Start a detached PESMaker process that submits a whole stage. + + Detaching the outer PESMaker process keeps a local ``bash`` submission + loop alive after its launching terminal or SSH connection is closed. + """ + output_dir = _stage_output_dir(config, stage) + output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + log_path = ( + output_dir / f"{stage}_submit_{timestamp}_{os.getpid()}.log" + ).resolve() + command = [ + sys.executable, + "-u", + "-m", + "pesmaker", + "submit", + str(config_path.resolve()), + "--stage", + stage, + ] + + with log_path.open("ab") as output: + common_options = { + "stdin": subprocess.DEVNULL, + "stdout": output, + "stderr": subprocess.STDOUT, + "close_fds": True, + } + if os.name == "nt": + process = subprocess.Popen( + command, + creationflags=( + subprocess.DETACHED_PROCESS + | subprocess.CREATE_NEW_PROCESS_GROUP + ), + **common_options, + ) + else: + process = subprocess.Popen( + command, + start_new_session=True, + **common_options, + ) + + return BackgroundSubmitProcess(pid=process.pid, log_path=log_path) + + def submit_jobs( config: PESMakerConfig, *, diff --git a/tests/test_submit_background.py b/tests/test_submit_background.py new file mode 100644 index 0000000..5555544 --- /dev/null +++ b/tests/test_submit_background.py @@ -0,0 +1,133 @@ +# Copyright 2026 Ting Liang and PESMaker development team +# This file is part of PESMaker. +"""Tests for detached whole-stage submission.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +from pesmaker.cli import main +from pesmaker.config.io import load_config +from pesmaker.jobs.submit import ( + BackgroundSubmitProcess, + start_background_submit, +) + + +def _write_config(path: Path, output_dir: Path) -> None: + path.write_text( + f"""project: background_submit_test +labeling: + engine: vasp + output_dir: {output_dir.as_posix()} +jobs: + submit_command: bash +""", + encoding="utf-8", + ) + + +def test_start_background_submit_detaches_outer_pesmaker_process( + tmp_path, monkeypatch +): + """The worker should have no terminal streams and use a new session.""" + output_dir = tmp_path / "labeling" + config_path = tmp_path / "run.yaml" + _write_config(config_path, output_dir) + calls = [] + + def fake_popen(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace(pid=4321) + + monkeypatch.setattr("pesmaker.jobs.submit.subprocess.Popen", fake_popen) + + result = start_background_submit( + load_config(config_path), config_path, stage="scf" + ) + + assert result.pid == 4321 + assert result.log_path.parent == output_dir + assert result.log_path.name.startswith("scf_submit_") + assert result.log_path.suffix == ".log" + assert result.log_path.exists() + command, kwargs = calls[0] + assert command == [ + sys.executable, + "-u", + "-m", + "pesmaker", + "submit", + str(config_path.resolve()), + "--stage", + "scf", + ] + assert kwargs["stdin"] is subprocess.DEVNULL + assert kwargs["stderr"] is subprocess.STDOUT + assert kwargs["close_fds"] is True + assert Path(kwargs["stdout"].name) == result.log_path + if os.name == "nt": + assert kwargs["creationflags"] == ( + subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP + ) + assert "start_new_session" not in kwargs + else: + assert kwargs["start_new_session"] is True + assert "creationflags" not in kwargs + + +def test_cli_submit_background_starts_worker_and_returns( + tmp_path, monkeypatch, capsys +): + """The foreground CLI should report the detached PID and log immediately.""" + output_dir = tmp_path / "labeling" + config_path = tmp_path / "run.yaml" + log_path = output_dir / "scf_submit_test.log" + _write_config(config_path, output_dir) + calls = [] + + def fake_start(config, path, *, stage): + calls.append((config.project, path, stage)) + return BackgroundSubmitProcess(pid=9876, log_path=log_path) + + monkeypatch.setattr("pesmaker.cli.start_background_submit", fake_start) + + assert main(["submit", str(config_path), "--background"]) == 0 + output = capsys.readouterr().out + + assert calls == [("background_submit_test", config_path, "scf")] + assert "Background submission started." in output + assert "Process ID : 9876" in output + assert f"Log : {log_path}" in output + + +def test_cli_rejects_background_dry_run(tmp_path, monkeypatch, capsys): + """A preview should stay foreground instead of spawning a detached worker.""" + config_path = tmp_path / "run.yaml" + _write_config(config_path, tmp_path / "labeling") + started = False + + def fake_start(*args, **kwargs): + nonlocal started + started = True + + monkeypatch.setattr("pesmaker.cli.start_background_submit", fake_start) + + assert ( + main( + [ + "submit", + str(config_path), + "--background", + "--dry-run", + ] + ) + == 2 + ) + + assert started is False + assert "--background cannot be used with --dry-run" in capsys.readouterr().err From af68fbc16450e44d4187c7a47ee717f5c90e238f Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Wed, 15 Jul 2026 17:14:45 +0800 Subject: [PATCH 3/8] =?UTF-8?q?Added=20a=20task=20progress=20indicator=20f?= =?UTF-8?q?or=20single-point=20calculations=20after=20=E2=80=98pesmaker=20?= =?UTF-8?q?submit=E2=80=98.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/commands/scf-setup.md | 41 ++++++++++++++ docs/commands/submit.md | 16 ++++++ src/pesmaker/jobs/submit.py | 97 ++++++++++++++++++++++++++++++--- tests/test_submit_background.py | 93 +++++++++++++++++++++++++++++++ 4 files changed, 239 insertions(+), 8 deletions(-) diff --git a/docs/commands/scf-setup.md b/docs/commands/scf-setup.md index 2026629..fee875b 100644 --- a/docs/commands/scf-setup.md +++ b/docs/commands/scf-setup.md @@ -83,6 +83,47 @@ PESMaker prepares one SCF job per frame. This lets users point `labeling.input_dir` at a manually prepared folder even when the folder was not created by PESMaker and has no `manifest.jsonl`. +## Distribute Common Files To Every SCF Folder + +Use `labeling.template_dir` to copy a set of common auxiliary files into every +prepared SCF calculation folder: + +```yaml +labeling: + engine: vasp + input_dir: generated + output_dir: labeling + template_dir: templates/vasp/common +``` + +For example, given this source directory: + +```text +templates/vasp/common/ + KPOINTS + vdW_kernel.bindat + helper.dat + nested/ + ignored.dat +``` + +each prepared calculation folder receives `KPOINTS`, `vdW_kernel.bindat`, and +`helper.dat`. PESMaker copies only regular files directly inside +`template_dir`; it does not recursively copy `nested/` or any other +subdirectory. File names are preserved, and existing destination files with +the same names are overwritten at the time of copying. File metadata is +preserved where supported by the operating system. + +Use this option for auxiliary files that should be identical in every SCF job. +Avoid placing PESMaker-managed names such as `POSCAR`, `INCAR`, `POTCAR`, or +`submit.sh` in `template_dir`, because setup ordering can cause one version to +overwrite another. In particular, PESMaker writes the final POTCAR and submit +script after distributing these common files. + +`labeling.template_dir` is separate from `jobs.copy_sub_file`. The former +distributes all top-level files from a directory, while the latter copies one +submit-script template specifically to `submit.sh`. + ## Outputs ```text diff --git a/docs/commands/submit.md b/docs/commands/submit.md index 9933913..a7f9b97 100644 --- a/docs/commands/submit.md +++ b/docs/commands/submit.md @@ -81,6 +81,22 @@ This option protects the outer PESMaker process. It is especially useful with serial sequence alive and starts the next calculation only after the current `submit.sh` exits. Closing the terminal does not stop that worker. +For synchronous `bash` submissions, PESMaker reports every calculation as it +starts and as soon as its `submit.sh` exits: + +```text +[2026-07-15 18:20:03] STARTED 1/120 labeling/calc_000000 +[2026-07-15 18:52:41] COMPLETED 1/120 labeling/calc_000000 elapsed=00:32:38 +[2026-07-15 18:52:41] STARTED 2/120 labeling/calc_000001 +``` + +A nonzero script exit is reported as `FAILED` before the serial sequence stops. +Each event is flushed immediately to both the terminal and +`labeling/scf_submitted_jobs.txt`. With `--background`, terminal output is +redirected to the printed `scf_submit_*.log`, so both logs can be monitored +while the calculations are running. `COMPLETED` means that `submit.sh` exited +successfully; it does not by itself certify VASP electronic convergence. + Monitor it with the printed log path or the normal calculation output files: ```bash diff --git a/src/pesmaker/jobs/submit.py b/src/pesmaker/jobs/submit.py index af04ed7..d9c89e1 100644 --- a/src/pesmaker/jobs/submit.py +++ b/src/pesmaker/jobs/submit.py @@ -24,6 +24,7 @@ import shlex import subprocess import sys +import time from pathlib import Path from pesmaker.artifacts import _read_manifest, _section_output_dir @@ -141,7 +142,20 @@ def submit_jobs( skip_completed = _skip_completed_jobs(config, stage) check_scf_convergence = _check_scf_convergence(config, stage) refresh_scripts = _refresh_vasp_submit_scripts(config, stage, skip_completed) - for job in jobs: + live_progress = _is_foreground_bash_submit(submit_command) and not dry_run + if not dry_run: + submitted_log.write_text("", encoding="utf-8") + + def record(line: str, *, echo: bool = False) -> None: + lines.append(line) + if not dry_run: + with submitted_log.open("a", encoding="utf-8") as handle: + handle.write(f"{line}\n") + handle.flush() + if echo: + print(line, flush=True) + + for job_number, job in enumerate(jobs, start=1): workdir = job.workdir if skip_completed: outcar_state = _vasp_outcar_state(workdir) @@ -149,16 +163,16 @@ def submit_jobs( outcar_state, check_scf_convergence=check_scf_convergence, ): - lines.append(f"SKIPPED completed VASP job: {workdir}") + record(f"SKIPPED completed VASP job: {workdir}") skipped_count += 1 continue retry_reason = _vasp_retry_reason(outcar_state) if retry_reason: - lines.append(f"RETRY {retry_reason}: {workdir}") + record(f"RETRY {retry_reason}: {workdir}") if refresh_scripts: script = _refresh_vasp_submit_script(config, workdir) - lines.append(f"REFRESHED submit script: {script}") + record(f"REFRESHED submit script: {script}") else: script = _existing_submit_script(job) if script is None: @@ -166,13 +180,46 @@ def submit_jobs( display = _submit_display(submit_command, script) if dry_run: - lines.append(f"DRY-RUN {display}") + record(f"DRY-RUN {display}") submitted_count += 1 continue - message = _run_submit_command(submit_command, script) - lines.append(f"{script.parent}: {message}") + + started_at = time.monotonic() + if live_progress: + record( + _progress_line("STARTED", job_number, len(jobs), workdir), + echo=True, + ) + try: + message = _run_submit_command(submit_command, script) + except (OSError, subprocess.SubprocessError): + if live_progress: + record( + _progress_line( + "FAILED", + job_number, + len(jobs), + workdir, + elapsed=time.monotonic() - started_at, + ), + echo=True, + ) + raise + if live_progress: + record( + _progress_line( + "COMPLETED", + job_number, + len(jobs), + workdir, + elapsed=time.monotonic() - started_at, + ), + echo=True, + ) + record(f"{script.parent}: {message}") submitted_count += 1 - submitted_log.write_text("\n".join(lines) + "\n", encoding="utf-8") + if dry_run: + submitted_log.write_text("\n".join(lines) + "\n", encoding="utf-8") action = "Would submit" if dry_run else "Submitted" message = f"{action} {submitted_count} {stage} job(s)" if skipped_count: @@ -346,6 +393,40 @@ def _is_nohup_submit(submit_command: str) -> bool: return shlex.split(submit_command) == ["nohup"] +def _is_foreground_bash_submit(submit_command: str) -> bool: + """Return whether submit scripts run synchronously through Bash.""" + command = shlex.split(submit_command) + if not command: + return False + return Path(command[0]).name.lower() in {"bash", "bash.exe"} + + +def _progress_line( + status: str, + job_number: int, + total_jobs: int, + workdir: Path, + *, + elapsed: float | None = None, +) -> str: + """Format one timestamped live progress event.""" + line = ( + f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] " + f"{status:<9} {job_number}/{total_jobs} {workdir}" + ) + if elapsed is not None: + line += f" elapsed={_format_elapsed(elapsed)}" + return line + + +def _format_elapsed(elapsed: float) -> str: + """Format elapsed seconds as HH:MM:SS without wrapping after 24 hours.""" + total_seconds = max(0, int(elapsed)) + hours, remainder = divmod(total_seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" + + def _stage_prepared_jobs( config: PESMakerConfig, stage: str, diff --git a/tests/test_submit_background.py b/tests/test_submit_background.py index 5555544..c0883c2 100644 --- a/tests/test_submit_background.py +++ b/tests/test_submit_background.py @@ -4,17 +4,21 @@ from __future__ import annotations +import json import os import subprocess import sys from pathlib import Path from types import SimpleNamespace +import pytest + from pesmaker.cli import main from pesmaker.config.io import load_config from pesmaker.jobs.submit import ( BackgroundSubmitProcess, start_background_submit, + submit_jobs, ) @@ -131,3 +135,92 @@ def fake_start(*args, **kwargs): assert started is False assert "--background cannot be used with --dry-run" in capsys.readouterr().err + + +def _write_prepared_jobs(output_dir: Path, count: int) -> list[Path]: + output_dir.mkdir(parents=True) + workdirs = [] + records = [] + for index in range(count): + workdir = output_dir / f"calc_{index:06d}" + workdir.mkdir() + script = workdir / "submit.sh" + script.write_text("#!/usr/bin/env bash\n", encoding="utf-8") + workdirs.append(workdir) + records.append( + { + "workdir": str(workdir), + "submit_script": str(script), + } + ) + (output_dir / "labeling_manifest.jsonl").write_text( + "\n".join(json.dumps(record) for record in records) + "\n", + encoding="utf-8", + ) + return workdirs + + +def test_bash_submit_reports_and_flushes_each_completed_job( + tmp_path, monkeypatch, capsys +): + """Each synchronous Bash job should be visible before the next one runs.""" + output_dir = tmp_path / "labeling" + workdirs = _write_prepared_jobs(output_dir, 2) + config_path = tmp_path / "run.yaml" + _write_config(config_path, output_dir) + monotonic_values = iter((10.0, 15.0, 20.0, 27.0)) + calls = [] + + def fake_run(submit_command, script): + calls.append(script.parent) + current_log = (output_dir / "scf_submitted_jobs.txt").read_text( + encoding="utf-8" + ) + assert f" STARTED {len(calls)}/2 {script.parent}" in current_log + if len(calls) == 2: + assert f" COMPLETED 1/2 {workdirs[0]}" in current_log + return f"finished {script.parent.name}" + + monkeypatch.setattr("pesmaker.jobs.submit.time.monotonic", lambda: next(monotonic_values)) + monkeypatch.setattr("pesmaker.jobs.submit._run_submit_command", fake_run) + + result = submit_jobs(load_config(config_path)) + output = capsys.readouterr().out + log_text = result.files[0].read_text(encoding="utf-8") + + assert calls == workdirs + assert f" STARTED 1/2 {workdirs[0]}" in output + assert f" COMPLETED 1/2 {workdirs[0]} elapsed=00:00:05" in output + assert f" STARTED 2/2 {workdirs[1]}" in output + assert f" COMPLETED 2/2 {workdirs[1]} elapsed=00:00:07" in output + assert f" STARTED 1/2 {workdirs[0]}" in log_text + assert f" COMPLETED 1/2 {workdirs[0]} elapsed=00:00:05" in log_text + assert log_text.index("COMPLETED 1/2") < log_text.index("STARTED 2/2") + + +def test_bash_submit_reports_failure_before_stopping( + tmp_path, monkeypatch, capsys +): + """A failed foreground script should be flushed before the error escapes.""" + output_dir = tmp_path / "labeling" + workdir = _write_prepared_jobs(output_dir, 1)[0] + config_path = tmp_path / "run.yaml" + _write_config(config_path, output_dir) + monotonic_values = iter((100.0, 103.0)) + + def fake_run(submit_command, script): + raise subprocess.CalledProcessError(7, ["bash", script.name]) + + monkeypatch.setattr("pesmaker.jobs.submit.time.monotonic", lambda: next(monotonic_values)) + monkeypatch.setattr("pesmaker.jobs.submit._run_submit_command", fake_run) + + with pytest.raises(subprocess.CalledProcessError): + submit_jobs(load_config(config_path)) + + output = capsys.readouterr().out + log_text = (output_dir / "scf_submitted_jobs.txt").read_text( + encoding="utf-8" + ) + expected = f" FAILED 1/1 {workdir} elapsed=00:00:03" + assert expected in output + assert expected in log_text From 4763208818a53b333aa46eeeac96afa240935a58 Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Wed, 15 Jul 2026 20:06:51 +0800 Subject: [PATCH 4/8] Friendly termination signal add explanation to the termination signal "subprocess.CalledProcessError: Command '['bash', 'submit.sh']' returned non-zero exit status 143." when a vasp single point process is killed --- docs/commands/submit.md | 15 ++++++++++ src/pesmaker/__main__.py | 2 +- src/pesmaker/cli.py | 36 +++++++++++++++++++++- src/pesmaker/jobs/submit.py | 53 ++++++++++++++++++++++++++++++++- tests/test_submit_background.py | 38 ++++++++++++++++++++++- 5 files changed, 140 insertions(+), 4 deletions(-) diff --git a/docs/commands/submit.md b/docs/commands/submit.md index a7f9b97..0486517 100644 --- a/docs/commands/submit.md +++ b/docs/commands/submit.md @@ -97,6 +97,21 @@ redirected to the printed `scf_submit_*.log`, so both logs can be monitored while the calculations are running. `COMPLETED` means that `submit.sh` exited successfully; it does not by itself certify VASP electronic convergence. +The CLI converts a failed script into a concise stop summary instead of a +Python traceback. For example, shell status `143` conventionally means that the +script received `SIGTERM`: + +```text +SCF serial submission stopped. +Failed job : labeling/calc_000004 +Exit status : 143 +Termination : SIGTERM +Remaining jobs : 115 not started +``` + +The same summary is written to `scf_submit_*.log` during a background run. The +failed event remains in `scf_submitted_jobs.txt`. + Monitor it with the printed log path or the normal calculation output files: ```bash diff --git a/src/pesmaker/__main__.py b/src/pesmaker/__main__.py index f8586ab..84acb9c 100644 --- a/src/pesmaker/__main__.py +++ b/src/pesmaker/__main__.py @@ -18,4 +18,4 @@ from pesmaker.cli import main if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/src/pesmaker/cli.py b/src/pesmaker/cli.py index a4f2bc0..e459d8a 100644 --- a/src/pesmaker/cli.py +++ b/src/pesmaker/cli.py @@ -18,6 +18,7 @@ from __future__ import annotations import argparse +import signal import shlex import sys from pathlib import Path @@ -38,7 +39,11 @@ format_generate_summary, generate_structures, ) -from pesmaker.jobs.submit import start_background_submit, submit_jobs +from pesmaker.jobs.submit import ( + SubmissionJobError, + start_background_submit, + submit_jobs, +) from pesmaker.labelers.vasp import setup_labeling from pesmaker.plot import available_plot_commands, run_plot_command from pesmaker.results import StageResult @@ -317,6 +322,9 @@ def main(argv: list[str] | None = None) -> int: submit_command=str(config.jobs.options.get("submit_command", "sbatch")), ) return 0 + except SubmissionJobError as exc: + print(_format_submission_job_error(exc), file=sys.stderr) + return 2 except (OSError, ValueError) as exc: print( _format_cli_error(args.command, args.config, config, exc), file=sys.stderr @@ -367,6 +375,32 @@ def _format_cli_error( return f"Error: {message}" +def _format_submission_job_error(exc: SubmissionJobError) -> str: + """Format a failed submit script without exposing a Python traceback.""" + mode = "serial submission" if exc.serial else "submission" + termination = _termination_name(exc.returncode) + lines = [ + f"{exc.stage.upper()} {mode} stopped.", + f"Failed job : {exc.workdir}", + f"Exit status : {exc.returncode}", + ] + if termination is not None: + lines.append(f"Termination : {termination}") + lines.append(f"Remaining jobs : {exc.remaining_jobs} not started") + return "\n".join(lines) + + +def _termination_name(returncode: int) -> str | None: + """Decode conventional shell signal exit statuses such as 143.""" + signal_number = -returncode if returncode < 0 else returncode - 128 + if signal_number <= 0: + return None + try: + return signal.Signals(signal_number).name + except ValueError: + return None + + def _write_starter_config(path: Path) -> int: """Write a minimal starter YAML configuration file. diff --git a/src/pesmaker/jobs/submit.py b/src/pesmaker/jobs/submit.py index d9c89e1..698ccd4 100644 --- a/src/pesmaker/jobs/submit.py +++ b/src/pesmaker/jobs/submit.py @@ -67,6 +67,37 @@ class BackgroundSubmitProcess: log_path: Path +class SubmissionJobError(subprocess.CalledProcessError): + """One prepared job failed while a stage was being submitted.""" + + def __init__( + self, + error: subprocess.CalledProcessError, + *, + stage: str, + workdir: Path, + job_number: int, + total_jobs: int, + serial: bool, + ) -> None: + super().__init__( + error.returncode, + error.cmd, + output=error.output, + stderr=error.stderr, + ) + self.stage = stage + self.workdir = workdir + self.job_number = job_number + self.total_jobs = total_jobs + self.serial = serial + + @property + def remaining_jobs(self) -> int: + """Return the number of later manifest entries not started here.""" + return max(0, self.total_jobs - self.job_number) + + def start_background_submit( config: PESMakerConfig, config_path: Path, @@ -192,7 +223,27 @@ def record(line: str, *, echo: bool = False) -> None: ) try: message = _run_submit_command(submit_command, script) - except (OSError, subprocess.SubprocessError): + except subprocess.CalledProcessError as exc: + if live_progress: + record( + _progress_line( + "FAILED", + job_number, + len(jobs), + workdir, + elapsed=time.monotonic() - started_at, + ), + echo=True, + ) + raise SubmissionJobError( + exc, + stage=stage, + workdir=workdir, + job_number=job_number, + total_jobs=len(jobs), + serial=live_progress, + ) from None + except OSError: if live_progress: record( _progress_line( diff --git a/tests/test_submit_background.py b/tests/test_submit_background.py index c0883c2..1f143a6 100644 --- a/tests/test_submit_background.py +++ b/tests/test_submit_background.py @@ -17,6 +17,7 @@ from pesmaker.config.io import load_config from pesmaker.jobs.submit import ( BackgroundSubmitProcess, + SubmissionJobError, start_background_submit, submit_jobs, ) @@ -214,7 +215,7 @@ def fake_run(submit_command, script): monkeypatch.setattr("pesmaker.jobs.submit.time.monotonic", lambda: next(monotonic_values)) monkeypatch.setattr("pesmaker.jobs.submit._run_submit_command", fake_run) - with pytest.raises(subprocess.CalledProcessError): + with pytest.raises(SubmissionJobError) as caught: submit_jobs(load_config(config_path)) output = capsys.readouterr().out @@ -224,3 +225,38 @@ def fake_run(submit_command, script): expected = f" FAILED 1/1 {workdir} elapsed=00:00:03" assert expected in output assert expected in log_text + assert caught.value.returncode == 7 + assert caught.value.workdir == workdir + assert caught.value.remaining_jobs == 0 + + +def test_cli_formats_killed_bash_job_without_traceback( + tmp_path, monkeypatch, capsys +): + """A signal-style shell status should become a concise stop summary.""" + output_dir = tmp_path / "labeling" + workdirs = _write_prepared_jobs(output_dir, 2) + config_path = tmp_path / "run.yaml" + _write_config(config_path, output_dir) + monotonic_values = iter((200.0, 205.0)) + + def fake_run(submit_command, script): + raise subprocess.CalledProcessError(143, ["bash", script.name]) + + monkeypatch.setattr( + "pesmaker.jobs.submit.time.monotonic", + lambda: next(monotonic_values), + ) + monkeypatch.setattr("pesmaker.jobs.submit._run_submit_command", fake_run) + + assert main(["submit", str(config_path)]) == 2 + captured = capsys.readouterr() + + assert f" FAILED 1/2 {workdirs[0]} elapsed=00:00:05" in captured.out + assert "SCF serial submission stopped." in captured.err + assert f"Failed job : {workdirs[0]}" in captured.err + assert "Exit status : 143" in captured.err + assert "Termination : SIGTERM" in captured.err + assert "Remaining jobs : 1 not started" in captured.err + assert "Traceback" not in captured.err + assert "CalledProcessError" not in captured.err From 78213426ff2923ea93f8f02e389cd32b4624ed4b Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Wed, 22 Jul 2026 20:18:39 +0800 Subject: [PATCH 5/8] Add VASP AIMD FPS to NEP toolkit --- PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py | 667 +++++++++++++++++++++++ tests/test_aimd_fps_to_nep_toolkit.py | 252 +++++++++ 2 files changed, 919 insertions(+) create mode 100644 PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py create mode 100644 tests/test_aimd_fps_to_nep_toolkit.py diff --git a/PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py b/PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py new file mode 100644 index 0000000..7d8a87a --- /dev/null +++ b/PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py @@ -0,0 +1,667 @@ +#!/usr/bin/env python3 +"""Create a compact NEP training set from one fixed-cell VASP AIMD run. + +The toolkit deliberately stays outside PESMaker's workflow implementation. It +uses the public ``python -m pesmaker select`` command for FPS on XDATCAR, then +streams ionic frames from OUTCAR and writes labels only for the selected steps. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import json +import logging +import math +from pathlib import Path +import re +import subprocess +import sys +from typing import Any, Iterable, TextIO + +import numpy as np + + +OUTPUT_DIR_NAME = "PESMakerToolkit_AIMD_FPS_to_NEP" +CELL_TOLERANCE_ANGSTROM = 1.0e-4 +POSITION_TOLERANCE_ANGSTROM = 1.0e-4 +REQUIRED_FPS_FILES = ( + "selected.xyz", + "manifest.jsonl", + "selection_features.npy", + "fps_selection.png", +) + + +class ToolkitError(RuntimeError): + """A user-facing toolkit failure.""" + + +@dataclass(frozen=True) +class SelectionEntry: + """One FPS-selected XDATCAR frame and its expected OUTCAR ionic step.""" + + fps_order: int + source_frame: int + ionic_step: int + atoms: Any + + +@dataclass(frozen=True) +class GeometryCheck: + """Maximum differences observed when matching one selected frame.""" + + max_cell_difference: float + max_position_difference: float + + +@dataclass +class ExtractionStats: + """Dataset-level statistics accumulated during streaming extraction.""" + + written: int = 0 + min_energy: float = math.inf + max_energy: float = -math.inf + max_force: float = 0.0 + max_abs_virial: float = 0.0 + max_cell_difference: float = 0.0 + max_position_difference: float = 0.0 + + def update( + self, + *, + energy: float, + forces: np.ndarray, + virial: np.ndarray, + geometry: GeometryCheck, + ) -> None: + self.written += 1 + self.min_energy = min(self.min_energy, energy) + self.max_energy = max(self.max_energy, energy) + self.max_force = max( + self.max_force, + float(np.max(np.linalg.norm(forces, axis=1))), + ) + self.max_abs_virial = max( + self.max_abs_virial, + float(np.max(np.abs(virial))), + ) + self.max_cell_difference = max( + self.max_cell_difference, + geometry.max_cell_difference, + ) + self.max_position_difference = max( + self.max_position_difference, + geometry.max_position_difference, + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Run PESMaker FPS on a VASP AIMD XDATCAR and stream matching " + "energy/force/virial labels from OUTCAR into NEP train.xyz." + ) + ) + parser.add_argument( + "--aimd-dir", + type=Path, + required=True, + help="VASP AIMD directory containing INCAR, XDATCAR, and OUTCAR.", + ) + parser.add_argument( + "--max-count", + type=int, + required=True, + help="Maximum number of structures selected by FPS (must be positive).", + ) + parser.add_argument( + "--min-distance", + type=float, + default=0.0, + help=( + "Minimum FPS distance in PESMaker simple-descriptor space. " + "Default: 0.0." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + try: + _validate_cli_values(args.max_count, args.min_distance) + run_toolkit( + aimd_dir=args.aimd_dir, + max_count=args.max_count, + min_distance=args.min_distance, + ) + except ToolkitError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 2 + except KeyboardInterrupt: + print("Error: interrupted by user", file=sys.stderr) + return 130 + return 0 + + +def run_toolkit(*, aimd_dir: Path, max_count: int, min_distance: float) -> Path: + """Run the complete XDATCAR -> FPS -> streamed OUTCAR conversion.""" + aimd_dir = aimd_dir.expanduser().resolve() + incar, xdatcar, outcar = _validate_inputs(aimd_dir) + output_dir = aimd_dir / OUTPUT_DIR_NAME + if output_dir.exists(): + raise ToolkitError( + f"output directory already exists; refusing to overwrite: {output_dir}" + ) + + output_dir.mkdir(parents=False, exist_ok=False) + logger = _configure_logger(output_dir / "toolkit.log") + logger.info("PESMaker AIMD FPS to NEP toolkit") + logger.info("AIMD directory : %s", aimd_dir) + logger.info("XDATCAR : %s", xdatcar) + logger.info("OUTCAR : %s", outcar) + logger.info("Output : %s", output_dir) + logger.info("max_count : %d", max_count) + logger.info("min_distance : %.16g", min_distance) + + try: + nblock = parse_nblock(incar) + logger.info("NBLOCK : %d", nblock) + config_path = output_dir / "pesmaker_fps.yaml" + _write_pesmaker_config( + config_path, + xdatcar=xdatcar, + output_dir=output_dir, + max_count=max_count, + min_distance=min_distance, + ) + _run_pesmaker_select(config_path, aimd_dir=aimd_dir, logger=logger) + _verify_fps_outputs(output_dir) + + selections = load_selections(output_dir, nblock=nblock) + logger.info("FPS selected : %d frame(s)", len(selections)) + stats = extract_train_xyz( + outcar=outcar, + selections=selections, + output_dir=output_dir, + logger=logger, + ) + _log_success(logger, stats, output_dir / "train.xyz") + return output_dir / "train.xyz" + except Exception as exc: + if isinstance(exc, ToolkitError): + logger.error("Toolkit failed: %s", exc) + raise + logger.exception("Toolkit failed with an unexpected error") + raise ToolkitError(str(exc)) from exc + + +def _validate_cli_values(max_count: int, min_distance: float) -> None: + if max_count < 1: + raise ToolkitError("--max-count must be a positive integer") + if not math.isfinite(min_distance) or min_distance < 0.0: + raise ToolkitError("--min-distance must be a finite non-negative number") + + +def _validate_inputs(aimd_dir: Path) -> tuple[Path, Path, Path]: + if not aimd_dir.is_dir(): + raise ToolkitError(f"AIMD directory does not exist: {aimd_dir}") + paths = tuple(aimd_dir / name for name in ("INCAR", "XDATCAR", "OUTCAR")) + missing = [path.name for path in paths if not path.is_file()] + if missing: + raise ToolkitError( + f"AIMD directory is missing required file(s): {', '.join(missing)}" + ) + return paths + + +def parse_nblock(incar: Path) -> int: + """Return the last active NBLOCK value in INCAR, defaulting to one.""" + nblock = 1 + pattern = re.compile(r"^\s*NBLOCK\s*=\s*([^\s;]+)", re.IGNORECASE) + for raw_line in incar.read_text(encoding="utf-8", errors="replace").splitlines(): + line = raw_line.split("!", 1)[0].split("#", 1)[0] + match = pattern.match(line) + if not match: + continue + try: + nblock = int(match.group(1)) + except ValueError as exc: + raise ToolkitError(f"invalid NBLOCK value in {incar}: {match.group(1)}") from exc + if nblock < 1: + raise ToolkitError(f"NBLOCK must be positive in {incar}, got {nblock}") + return nblock + + +def _configure_logger(path: Path) -> logging.Logger: + logger = logging.getLogger(f"pesmaker_aimd_toolkit.{id(path)}") + logger.setLevel(logging.INFO) + logger.propagate = False + formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s") + file_handler = logging.FileHandler(path, encoding="utf-8") + file_handler.setFormatter(formatter) + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(logging.Formatter("%(message)s")) + logger.addHandler(file_handler) + logger.addHandler(stream_handler) + return logger + + +def _write_pesmaker_config( + path: Path, + *, + xdatcar: Path, + output_dir: Path, + max_count: int, + min_distance: float, +) -> None: + try: + import yaml + except ImportError as exc: # pragma: no cover - PESMaker requires PyYAML + raise ToolkitError("PyYAML is required to generate the FPS config") from exc + config = { + "project": "aimd_fps_to_nep", + "sampling": { + "engine": "none", + "selection": { + "method": "fps", + "trajectory_pattern": str(xdatcar), + "trajectory_format": "vasp-xdatcar", + "output_dir": str(output_dir), + "max_count": max_count, + "min_distance": min_distance, + "separate_trajectories": False, + "plot": True, + }, + }, + } + path.write_text( + yaml.safe_dump(config, sort_keys=False, allow_unicode=True), + encoding="utf-8", + ) + + +def _run_pesmaker_select( + config_path: Path, + *, + aimd_dir: Path, + logger: logging.Logger, +) -> None: + command = [sys.executable, "-m", "pesmaker", "select", str(config_path)] + logger.info("Running : %s", " ".join(command)) + try: + process = subprocess.Popen( + command, + cwd=aimd_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + ) + except OSError as exc: + raise ToolkitError(f"could not start PESMaker: {exc}") from exc + assert process.stdout is not None + try: + for line in process.stdout: + logger.info("[pesmaker] %s", line.rstrip()) + returncode = process.wait() + except KeyboardInterrupt: + process.terminate() + process.wait() + raise + if returncode != 0: + raise ToolkitError(f"PESMaker select failed with exit status {returncode}") + + +def _verify_fps_outputs(output_dir: Path) -> None: + missing = [name for name in REQUIRED_FPS_FILES if not (output_dir / name).is_file()] + if missing: + raise ToolkitError( + "PESMaker select completed without required output file(s): " + + ", ".join(missing) + ) + + +def load_selections(output_dir: Path, *, nblock: int) -> list[SelectionEntry]: + """Load the FPS manifest and selected frames in their common order.""" + manifest_path = output_dir / "manifest.jsonl" + records = [] + for line_number, line in enumerate( + manifest_path.read_text(encoding="utf-8").splitlines(), start=1 + ): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as exc: + raise ToolkitError( + f"invalid JSON in {manifest_path} line {line_number}" + ) from exc + records.append(record) + + try: + from ase.io import iread + except ImportError as exc: # pragma: no cover - PESMaker requires ASE + raise ToolkitError("ASE is required to read selected.xyz") from exc + selected_frames = list( + iread(output_dir / "selected.xyz", index=":", format="extxyz") + ) + if not records: + raise ToolkitError("FPS manifest contains no selected frames") + if len(records) != len(selected_frames): + raise ToolkitError( + "FPS manifest/selected.xyz frame count mismatch: " + f"{len(records)} vs {len(selected_frames)}" + ) + _validate_selected_trajectory_assumptions(selected_frames) + + selections = [] + seen_source_frames: set[int] = set() + seen_ionic_steps: set[int] = set() + for order, (record, atoms) in enumerate(zip(records, selected_frames)): + if "source_frame" not in record: + raise ToolkitError(f"manifest record {order} has no source_frame") + try: + source_frame = int(record["source_frame"]) + except (TypeError, ValueError) as exc: + raise ToolkitError( + f"manifest record {order} has invalid source_frame" + ) from exc + if source_frame < 0 or source_frame in seen_source_frames: + raise ToolkitError( + f"manifest contains invalid or duplicate source_frame: {source_frame}" + ) + ionic_step = (source_frame + 1) * nblock + if ionic_step in seen_ionic_steps: + raise ToolkitError(f"duplicate mapped ionic step: {ionic_step}") + seen_source_frames.add(source_frame) + seen_ionic_steps.add(ionic_step) + selections.append( + SelectionEntry( + fps_order=order, + source_frame=source_frame, + ionic_step=ionic_step, + atoms=atoms, + ) + ) + return selections + + +def _validate_selected_trajectory_assumptions(selected_frames: list[Any]) -> None: + """Enforce the first version's fixed-cell, fixed-order scope.""" + reference = selected_frames[0] + reference_symbols = reference.get_chemical_symbols() + reference_cell = np.asarray(reference.cell.array, dtype=float) + if reference_cell.shape != (3, 3) or not np.all(np.isfinite(reference_cell)): + raise ToolkitError("selected XDATCAR reference cell is not a finite 3x3 matrix") + + for frame_index, atoms in enumerate(selected_frames[1:], start=1): + if atoms.get_chemical_symbols() != reference_symbols: + raise ToolkitError( + "selected XDATCAR frames do not keep a fixed atom count/order " + f"(selected frame {frame_index})" + ) + cell = np.asarray(atoms.cell.array, dtype=float) + if cell.shape != (3, 3) or not np.all(np.isfinite(cell)): + raise ToolkitError( + f"selected XDATCAR frame {frame_index} cell is not a finite 3x3 matrix" + ) + max_difference = float(np.max(np.abs(cell - reference_cell))) + if max_difference > CELL_TOLERANCE_ANGSTROM: + raise ToolkitError( + "selected XDATCAR is not a fixed-cell trajectory: maximum cell " + f"difference {max_difference:.6g} A exceeds " + f"{CELL_TOLERANCE_ANGSTROM:g} A at selected frame {frame_index}" + ) + + +def extract_train_xyz( + *, + outcar: Path, + selections: list[SelectionEntry], + output_dir: Path, + logger: logging.Logger, +) -> ExtractionStats: + """Stream OUTCAR and write only the ionic steps selected by FPS.""" + by_step = {selection.ionic_step: selection for selection in selections} + wanted_steps = set(by_step) + found_steps: set[int] = set() + train_partial = output_dir / "train.xyz.partial" + mapping_partial = output_dir / "frame_mapping.jsonl.partial" + train_path = output_dir / "train.xyz" + mapping_path = output_dir / "frame_mapping.jsonl" + stats = ExtractionStats() + + logger.info("Streaming OUTCAR labels; the file is not loaded into memory.") + try: + frames = _iter_outcar_frames(outcar) + with train_partial.open("x", encoding="utf-8", newline="\n") as train_handle, \ + mapping_partial.open("x", encoding="utf-8", newline="\n") as map_handle: + for ionic_step, labeled_atoms in enumerate(frames, start=1): + selection = by_step.get(ionic_step) + if selection is None: + continue + geometry = validate_geometry(selection.atoms, labeled_atoms) + energy, forces, virial = extract_labels(labeled_atoms) + _write_nep_frame( + train_handle, + atoms=labeled_atoms, + energy=energy, + forces=forces, + virial=virial, + ) + mapping = { + "train_frame": stats.written, + "fps_order": selection.fps_order, + "source_frame": selection.source_frame, + "ionic_step": ionic_step, + "label_source": str(outcar), + } + map_handle.write(json.dumps(mapping, ensure_ascii=False) + "\n") + found_steps.add(ionic_step) + stats.update( + energy=energy, + forces=forces, + virial=virial, + geometry=geometry, + ) + logger.info( + "Matched ionic step %d (%d/%d)", + ionic_step, + len(found_steps), + len(wanted_steps), + ) + if found_steps == wanted_steps: + break + except ToolkitError: + raise + except Exception as exc: + raise ToolkitError(f"failed while streaming OUTCAR: {exc}") from exc + + missing = sorted(wanted_steps - found_steps) + if missing: + preview = ", ".join(str(step) for step in missing[:10]) + suffix = "..." if len(missing) > 10 else "" + raise ToolkitError( + f"OUTCAR did not provide {len(missing)} selected ionic step(s): " + f"{preview}{suffix}" + ) + + train_partial.replace(train_path) + mapping_partial.replace(mapping_path) + return stats + + +def _iter_outcar_frames(outcar: Path) -> Iterable[Any]: + """Yield ASE Atoms while parsing only labels required by this toolkit. + + ASE's default OUTCAR reader also parses k-point metadata. Some valid VASP + AIMD OUTCAR files contain abbreviated k-point header blocks that make that + unrelated parser fail before the first ionic frame. Building ASE's stream + generator with the minimal header/chunk parser set avoids that failure and + still uses ASE's native VASP unit and sign conversions. + """ + try: + from ase.io.vasp_parsers.vasp_outcar_parsers import ( + Cell, + Energy, + IonsPerSpecies, + OutcarChunkParser, + OutcarHeaderParser, + PositionsAndForces, + SpeciesTypes, + Spinpol, + Stress, + outcarchunks, + ) + except ImportError as exc: # pragma: no cover - PESMaker requires ASE + raise ToolkitError( + "ASE with vasp_outcar_parsers support is required to stream OUTCAR" + ) from exc + + header_parser = OutcarHeaderParser( + parsers=[SpeciesTypes(), IonsPerSpecies(), Spinpol()] + ) + chunk_parser = OutcarChunkParser( + parsers=[Cell(), PositionsAndForces(), Stress(), Energy()] + ) + with outcar.open("r", encoding="utf-8", errors="replace") as handle: + for chunk in outcarchunks( + handle, + header_parser=header_parser, + chunk_parser=chunk_parser, + ): + yield chunk.build() + + +def validate_geometry(selected_atoms: Any, labeled_atoms: Any) -> GeometryCheck: + """Verify that an XDATCAR selection matches one OUTCAR ionic frame.""" + if len(selected_atoms) != len(labeled_atoms): + raise ToolkitError( + "selected/OUTCAR atom-count mismatch: " + f"{len(selected_atoms)} vs {len(labeled_atoms)}" + ) + if selected_atoms.get_chemical_symbols() != labeled_atoms.get_chemical_symbols(): + raise ToolkitError("selected/OUTCAR element order mismatch") + + selected_cell = np.asarray(selected_atoms.cell.array, dtype=float) + labeled_cell = np.asarray(labeled_atoms.cell.array, dtype=float) + if selected_cell.shape != (3, 3) or labeled_cell.shape != (3, 3): + raise ToolkitError("selected/OUTCAR cell is not a 3x3 matrix") + if not np.all(np.isfinite(selected_cell)) or not np.all(np.isfinite(labeled_cell)): + raise ToolkitError("selected/OUTCAR cell contains non-finite values") + max_cell_difference = float(np.max(np.abs(selected_cell - labeled_cell))) + if max_cell_difference > CELL_TOLERANCE_ANGSTROM: + raise ToolkitError( + "selected/OUTCAR fixed-cell mismatch: maximum cell difference " + f"{max_cell_difference:.6g} A exceeds {CELL_TOLERANCE_ANGSTROM:g} A" + ) + + try: + selected_scaled = np.asarray( + selected_atoms.get_scaled_positions(wrap=False), dtype=float + ) + labeled_scaled = np.asarray( + labeled_atoms.get_scaled_positions(wrap=False), dtype=float + ) + except (ValueError, np.linalg.LinAlgError) as exc: + raise ToolkitError(f"could not calculate fractional coordinates: {exc}") from exc + delta_scaled = labeled_scaled - selected_scaled + delta_scaled -= np.round(delta_scaled) + delta_cartesian = delta_scaled @ selected_cell + max_position_difference = float( + np.max(np.linalg.norm(delta_cartesian, axis=1)) + ) + if max_position_difference > POSITION_TOLERANCE_ANGSTROM: + raise ToolkitError( + "selected/OUTCAR position mismatch: maximum minimum-image " + f"difference {max_position_difference:.6g} A exceeds " + f"{POSITION_TOLERANCE_ANGSTROM:g} A" + ) + return GeometryCheck(max_cell_difference, max_position_difference) + + +def extract_labels(atoms: Any) -> tuple[float, np.ndarray, np.ndarray]: + """Extract force-consistent energy, forces, and GPUMD-sign virial.""" + if atoms.calc is None: + raise ToolkitError("OUTCAR frame has no ASE calculator results") + results = atoms.calc.results + if "free_energy" not in results or results["free_energy"] is None: + raise ToolkitError("OUTCAR frame has no free_energy/TOTEN label") + energy = float(results["free_energy"]) + if not math.isfinite(energy): + raise ToolkitError("OUTCAR frame has a non-finite free_energy/TOTEN") + + if "forces" not in results or results["forces"] is None: + raise ToolkitError("OUTCAR frame has no force labels") + forces = np.asarray(results["forces"], dtype=float) + if forces.shape != (len(atoms), 3): + raise ToolkitError( + f"OUTCAR force shape is {forces.shape}, expected ({len(atoms)}, 3)" + ) + if not np.all(np.isfinite(forces)): + raise ToolkitError("OUTCAR frame has non-finite force labels") + + if "stress" not in results or results["stress"] is None: + raise ToolkitError("OUTCAR frame has no stress label") + try: + stress = np.asarray(atoms.get_stress(voigt=False), dtype=float) + except Exception as exc: + raise ToolkitError(f"could not obtain 3x3 ASE stress: {exc}") from exc + if stress.shape != (3, 3) or not np.all(np.isfinite(stress)): + raise ToolkitError("OUTCAR stress is not a finite 3x3 tensor") + volume = float(atoms.get_volume()) + if not math.isfinite(volume) or volume <= 0.0: + raise ToolkitError(f"OUTCAR frame has invalid cell volume: {volume}") + virial = -volume * stress + if not np.all(np.isfinite(virial)): + raise ToolkitError("converted virial contains non-finite values") + return energy, forces, virial + + +def _write_nep_frame( + handle: TextIO, + *, + atoms: Any, + energy: float, + forces: np.ndarray, + virial: np.ndarray, +) -> None: + cell = np.asarray(atoms.cell.array, dtype=float) + positions = np.asarray(atoms.positions, dtype=float) + lattice_text = " ".join(_number(value) for value in cell.reshape(-1)) + virial_text = " ".join(_number(value) for value in virial.reshape(-1)) + handle.write(f"{len(atoms)}\n") + handle.write( + f'Lattice="{lattice_text}" Energy={_number(energy)} ' + "Properties=species:S:1:pos:R:3:force:R:3 " + f'Virial="{virial_text}" pbc="T T T"\n' + ) + for symbol, position, force in zip( + atoms.get_chemical_symbols(), positions, forces + ): + values = " ".join(_number(value) for value in np.concatenate((position, force))) + handle.write(f"{symbol} {values}\n") + + +def _number(value: float) -> str: + return f"{float(value):.16g}" + + +def _log_success(logger: logging.Logger, stats: ExtractionStats, train_path: Path) -> None: + logger.info("Dataset complete") + logger.info("Structures : %d", stats.written) + logger.info("Energy range : %.16g .. %.16g eV", stats.min_energy, stats.max_energy) + logger.info("Maximum force : %.16g eV/A", stats.max_force) + logger.info("Maximum |Virial| : %.16g eV", stats.max_abs_virial) + logger.info("Maximum cell diff: %.16g A", stats.max_cell_difference) + logger.info("Maximum pos diff : %.16g A", stats.max_position_difference) + logger.info("train.xyz : %s", train_path) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_aimd_fps_to_nep_toolkit.py b/tests/test_aimd_fps_to_nep_toolkit.py new file mode 100644 index 0000000..5332621 --- /dev/null +++ b/tests/test_aimd_fps_to_nep_toolkit.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import importlib.util +import json +import logging +from pathlib import Path +import sys + +import numpy as np +import pytest +import yaml +from ase import Atoms +from ase.calculators.singlepoint import SinglePointCalculator +from ase.io import read, write + + +TOOL_PATH = ( + Path(__file__).resolve().parents[1] + / "PESMaker_AIMD_Toolkit" + / "aimd_fps_to_nep.py" +) +SPEC = importlib.util.spec_from_file_location("aimd_fps_to_nep", TOOL_PATH) +assert SPEC is not None and SPEC.loader is not None +tool = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = tool +SPEC.loader.exec_module(tool) + + +def test_parse_nblock_defaults_and_uses_last_active_value(tmp_path: Path): + incar = tmp_path / "INCAR" + incar.write_text( + "NBLOCK = 2 ! old setting\n# NBLOCK = 50\n nblock = 5\n", + encoding="utf-8", + ) + assert tool.parse_nblock(incar) == 5 + + +def test_parse_nblock_defaults_to_one(tmp_path: Path): + incar = tmp_path / "INCAR" + incar.write_text("ENCUT = 500\n", encoding="utf-8") + assert tool.parse_nblock(incar) == 1 + + +def test_validate_geometry_accepts_periodic_wrap(): + selected = Atoms( + "H", + scaled_positions=[[0.99999, 0.2, 0.3]], + cell=np.eye(3) * 10.0, + pbc=True, + ) + labeled = Atoms( + "H", + scaled_positions=[[-0.00001, 0.2, 0.3]], + cell=np.eye(3) * 10.0, + pbc=True, + ) + result = tool.validate_geometry(selected, labeled) + assert result.max_position_difference < 1.0e-10 + + +def test_validate_geometry_rejects_atom_order_mismatch(): + selected = Atoms("HO", positions=[[0, 0, 0], [1, 0, 0]], cell=np.eye(3) * 5) + labeled = Atoms("OH", positions=[[0, 0, 0], [1, 0, 0]], cell=np.eye(3) * 5) + with pytest.raises(tool.ToolkitError, match="element order"): + tool.validate_geometry(selected, labeled) + + +def test_extract_labels_converts_stress_to_gpumd_virial(): + atoms = Atoms("H", positions=[[0, 0, 0]], cell=np.eye(3) * 2.0, pbc=True) + stress = np.array([1.0, 2.0, 3.0, 0.4, 0.5, 0.6]) + atoms.calc = SinglePointCalculator( + atoms, + energy=-1.9, + free_energy=-2.0, + forces=np.array([[0.1, 0.2, 0.3]]), + stress=stress, + ) + energy, forces, virial = tool.extract_labels(atoms) + assert energy == pytest.approx(-2.0) + assert forces.shape == (1, 3) + assert np.allclose(virial, -8.0 * atoms.get_stress(voigt=False)) + + +def test_load_selections_maps_source_frame_with_nblock(tmp_path: Path): + output = tmp_path + frames = [ + Atoms("H", positions=[[0, 0, 0]], cell=np.eye(3), pbc=True), + Atoms("H", positions=[[0.1, 0, 0]], cell=np.eye(3), pbc=True), + ] + write(output / "selected.xyz", frames, format="extxyz") + records = [{"source_frame": 1}, {"source_frame": 4}] + (output / "manifest.jsonl").write_text( + "".join(json.dumps(record) + "\n" for record in records), + encoding="utf-8", + ) + selections = tool.load_selections(output, nblock=5) + assert [selection.ionic_step for selection in selections] == [10, 25] + + +def test_load_selections_rejects_variable_cell(tmp_path: Path): + frames = [ + Atoms("H", positions=[[0, 0, 0]], cell=np.eye(3), pbc=True), + Atoms("H", positions=[[0, 0, 0]], cell=np.eye(3) * 1.01, pbc=True), + ] + write(tmp_path / "selected.xyz", frames, format="extxyz") + (tmp_path / "manifest.jsonl").write_text( + '{"source_frame": 0}\n{"source_frame": 1}\n', + encoding="utf-8", + ) + with pytest.raises(tool.ToolkitError, match="not a fixed-cell trajectory"): + tool.load_selections(tmp_path, nblock=1) + + +def test_nep_frame_is_ase_readable_and_contains_required_labels(tmp_path: Path): + atoms = Atoms("H", positions=[[0.1, 0.2, 0.3]], cell=np.eye(3) * 2, pbc=True) + forces = np.array([[0.4, 0.5, 0.6]]) + virial = np.arange(9, dtype=float).reshape(3, 3) + path = tmp_path / "train.xyz" + with path.open("w", encoding="utf-8") as handle: + tool._write_nep_frame( + handle, + atoms=atoms, + energy=-1.25, + forces=forces, + virial=virial, + ) + loaded = read(path, format="extxyz") + assert loaded.info["Energy"] == pytest.approx(-1.25) + assert np.allclose(loaded.arrays["force"], forces) + assert np.allclose(np.asarray(loaded.info["Virial"]).reshape(3, 3), virial) + + +def test_cli_value_validation(): + with pytest.raises(tool.ToolkitError, match="max-count"): + tool._validate_cli_values(0, 0.0) + with pytest.raises(tool.ToolkitError, match="min-distance"): + tool._validate_cli_values(1, -0.1) + + +def test_generated_fps_yaml_uses_absolute_paths(tmp_path: Path): + xdatcar = (tmp_path / "XDATCAR").resolve() + output_dir = (tmp_path / tool.OUTPUT_DIR_NAME).resolve() + output_dir.mkdir() + config_path = output_dir / "pesmaker_fps.yaml" + tool._write_pesmaker_config( + config_path, + xdatcar=xdatcar, + output_dir=output_dir, + max_count=25, + min_distance=0.0, + ) + + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + assert config["sampling"]["engine"] == "none" + selection = config["sampling"]["selection"] + assert selection == { + "method": "fps", + "trajectory_pattern": str(xdatcar), + "trajectory_format": "vasp-xdatcar", + "output_dir": str(output_dir), + "max_count": 25, + "min_distance": 0.0, + "separate_trajectories": False, + "plot": True, + } + + +def test_existing_output_directory_is_refused(tmp_path: Path): + for name in ("INCAR", "XDATCAR", "OUTCAR"): + (tmp_path / name).write_text("", encoding="utf-8") + (tmp_path / tool.OUTPUT_DIR_NAME).mkdir() + with pytest.raises(tool.ToolkitError, match="refusing to overwrite"): + tool.run_toolkit(aimd_dir=tmp_path, max_count=1, min_distance=0.0) + + +def _labeled_hydrogen(position: float, energy: float) -> Atoms: + atoms = Atoms( + "H", + positions=[[position, 0, 0]], + cell=np.eye(3) * 5.0, + pbc=True, + ) + atoms.calc = SinglePointCalculator( + atoms, + energy=energy, + free_energy=energy, + forces=np.array([[position, 0.0, 0.0]]), + stress=np.zeros(6), + ) + return atoms + + +def test_extract_train_xyz_streams_only_selected_steps(tmp_path: Path, monkeypatch): + frames = [ + _labeled_hydrogen(0.0, -1.0), + _labeled_hydrogen(0.1, -1.1), + _labeled_hydrogen(0.2, -1.2), + ] + selections = [ + tool.SelectionEntry(0, 0, 1, frames[0].copy()), + tool.SelectionEntry(1, 2, 3, frames[2].copy()), + ] + outcar = tmp_path / "OUTCAR" + outcar.touch() + + def fake_outcar_frames(path): + assert Path(path) == outcar + yield from frames + + monkeypatch.setattr(tool, "_iter_outcar_frames", fake_outcar_frames) + stats = tool.extract_train_xyz( + outcar=outcar, + selections=selections, + output_dir=tmp_path, + logger=logging.getLogger("test-streaming-success"), + ) + + assert stats.written == 2 + assert (tmp_path / "train.xyz").is_file() + assert not (tmp_path / "train.xyz.partial").exists() + mappings = [ + json.loads(line) + for line in (tmp_path / "frame_mapping.jsonl").read_text( + encoding="utf-8" + ).splitlines() + ] + assert [record["ionic_step"] for record in mappings] == [1, 3] + assert [record["source_frame"] for record in mappings] == [0, 2] + + +def test_extract_failure_keeps_partial_but_not_train_xyz(tmp_path: Path, monkeypatch): + atoms = Atoms("H", positions=[[0, 0, 0]], cell=np.eye(3) * 5.0, pbc=True) + atoms.calc = SinglePointCalculator( + atoms, + energy=-1.0, + free_energy=-1.0, + forces=np.zeros((1, 3)), + ) + selection = tool.SelectionEntry(0, 0, 1, atoms.copy()) + outcar = tmp_path / "OUTCAR" + outcar.touch() + monkeypatch.setattr(tool, "_iter_outcar_frames", lambda path: iter([atoms])) + + with pytest.raises(tool.ToolkitError, match="no stress label"): + tool.extract_train_xyz( + outcar=outcar, + selections=[selection], + output_dir=tmp_path, + logger=logging.getLogger("test-streaming-failure"), + ) + assert not (tmp_path / "train.xyz").exists() + assert (tmp_path / "train.xyz.partial").exists() From 8fd336a96ed98dba48cfa026332975ceaa7a6905 Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Thu, 23 Jul 2026 12:17:53 +0800 Subject: [PATCH 6/8] Add VASP AIMD interval to NEP toolkit --- PESMaker_AIMD_Toolkit/README.md | 369 +++++++++++++++++ PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py | 386 ++++++++++++++++++ tests/test_aimd_interval_to_nep_toolkit.py | 131 ++++++ 3 files changed, 886 insertions(+) create mode 100644 PESMaker_AIMD_Toolkit/README.md create mode 100644 PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py create mode 100644 tests/test_aimd_interval_to_nep_toolkit.py diff --git a/PESMaker_AIMD_Toolkit/README.md b/PESMaker_AIMD_Toolkit/README.md new file mode 100644 index 0000000..912aa24 --- /dev/null +++ b/PESMaker_AIMD_Toolkit/README.md @@ -0,0 +1,369 @@ +# PESMaker VASP AIMD → NEP Toolkit + +本目录提供两个独立于 PESMaker 工作流的辅助脚本，用于从单个 VASP AIMD 目录生成 GPUMD/NEP 可用的 `train.xyz`。 + +| 脚本 | 选择方式 | 适合用途 | +| --- | --- | --- | +| `aimd_fps_to_nep.py` | PESMaker 最远点采样（FPS） | 根据构型差异选择代表性结构 | +| `aimd_interval_to_nep.py` | 等间距采样 | 按固定轨迹帧间隔抽稀 AIMD 时间序列 | + +两个脚本都不会修改或导入 PESMaker 私有源码，而是调用正式命令： + +```text +python -m pesmaker select <自动生成的 YAML> +``` + +## 共同工作流程 + +```text +XDATCAR + → PESMaker 选择构型 + → 根据 manifest.jsonl 和 NBLOCK 定位 OUTCAR 离子步 + → 流式提取 TOTEN、力和 stress + → 构型匹配与标签检查 + → train.xyz +``` + +OUTCAR 一次只解析一个离子步，不会一次性把全部帧读入内存。FPS 仍需要 PESMaker 读取 XDATCAR 并计算几何特征；等间距脚本会先以文本流方式统计 XDATCAR 帧数。 + +## 适用范围 + +当前版本支持： + +- 单个 VASP AIMD 目录。 +- 固定晶胞轨迹。 +- AIMD 过程中原子数和原子顺序保持不变。 +- XDATCAR 与 OUTCAR 来自同一次计算。 +- 没有人工拼接的 XDATCAR 或 OUTCAR。 +- 三维周期性体系。 + +当前版本不处理： + +- 变胞 AIMD。 +- 多个 AIMD 目录合并。 +- 拼接或混用的续算轨迹。 +- train/test 划分。 +- 电子步收敛质量过滤。 +- 已有输出目录覆盖或续跑。 + +## 环境要求 + +必须使用已经安装 PESMaker 的同一个 Python 环境运行脚本。环境至少需要： + +- Python 3.10 或更高版本。 +- PESMaker。 +- ASE。 +- NumPy。 +- PyYAML。 +- Matplotlib（FPS 绘图需要）。 + +可以使用以下命令检查环境： + +```bash +python -m pesmaker --help +python -c "import ase, numpy, yaml; print(ase.__version__)" +``` + +脚本通过 `sys.executable -m pesmaker select` 启动 PESMaker。因此，运行脚本的 Python 环境必须就是安装 PESMaker 的环境。 + +## 预期输入 + +`--aimd-dir` 必须明确指向包含以下文件的目录： + +```text +VASP_AIMD/ +├── INCAR +├── XDATCAR +└── OUTCAR +``` + +三个 VASP 文件只读打开。脚本不会修改、移动、删除、压缩或重命名它们。 + +## FPS 脚本用法 + +```bash +python PESMaker_AIMD_Toolkit/aimd_fps_to_nep.py \ + --aimd-dir /absolute/path/to/VASP_AIMD \ + --max-count 200 \ + --min-distance 0.0 +``` + +参数： + +| 参数 | 是否必填 | 说明 | +| --- | --- | --- | +| `--aimd-dir` | 是 | 包含 INCAR、XDATCAR 和 OUTCAR 的目录 | +| `--max-count` | 是 | FPS 最多选择多少个构型，必须是正整数 | +| `--min-distance` | 否 | PESMaker 简单几何描述符空间中的最小距离，默认 `0.0`，不是 Å | + +首次使用建议保持 `--min-distance 0.0`，仅通过 `--max-count` 控制数量。 + +FPS 输出目录固定为： + +```text +/PESMakerToolkit_AIMD_FPS_to_NEP/ +``` + +成功后的预期输出： + +```text +PESMakerToolkit_AIMD_FPS_to_NEP/ +├── train.xyz +├── selected.xyz +├── manifest.jsonl +├── selection_features.npy +├── fps_selection.png +├── frame_mapping.jsonl +├── pesmaker_fps.yaml +└── toolkit.log +``` + +## 等间距脚本用法 + +最简调用： + +```bash +python PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py \ + --aimd-dir /absolute/path/to/VASP_AIMD \ + --interval 100 +``` + +指定起始帧和结束帧： + +```bash +python PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py \ + --aimd-dir /absolute/path/to/VASP_AIMD \ + --interval 100 \ + --start-frame 1000 \ + --end-frame 10000 +``` + +参数： + +| 参数 | 是否必填 | 默认值 | 说明 | +| --- | --- | --- | --- | +| `--aimd-dir` | 是 | 无 | 包含 INCAR、XDATCAR 和 OUTCAR 的目录 | +| `--interval` | 是 | 无 | 每隔多少个 XDATCAR 帧选择一个构型，必须是正整数 | +| `--start-frame` | 否 | `0` | 允许选择的第一个 XDATCAR 帧，采用零基编号 | +| `--end-frame` | 否 | 最后一帧 | 允许选择的最后一个 XDATCAR 帧，采用包含式边界 | + +### 帧号和结束边界 + +`start-frame=0` 表示 XDATCAR 中保存的第一帧，不表示 VASP 离子步 0。 + +例如： + +```text +--start-frame 100 +--end-frame 1000 +--interval 200 +``` + +实际选择： + +```text +100, 300, 500, 700, 900 +``` + +`end-frame` 是允许范围的包含式上边界，但只有落在等间距序列上的帧才会被选择。脚本不会为了包含结束帧而额外加入一个非等间距构型。 + +如果不设置 `--end-frame`，采样范围延伸到 XDATCAR 最后一帧。显式结束帧超过轨迹范围时会报错，不会静默截断。 + +### interval 与 NBLOCK + +`--interval` 的单位是 XDATCAR 帧，不是 VASP 离子步。候选 OUTCAR 离子步按以下公式计算： + +```text +ionic_step = (source_frame + 1) × NBLOCK +``` + +例如： + +```text +NBLOCK = 5 +--interval 100 +``` + +相邻训练构型对应的 VASP 离子步间隔约为： + +```text +100 × 5 = 500 +``` + +帧号公式只负责候选定位；每个候选仍必须通过晶胞、元素顺序和周期性坐标比较。 + +### 等间距脚本预期输出 + +输出目录固定为： + +```text +/PESMakerToolkit_AIMD_Interval_to_NEP/ +``` + +成功后的预期输出： + +```text +PESMakerToolkit_AIMD_Interval_to_NEP/ +├── train.xyz +├── selected.xyz +├── manifest.jsonl +├── frame_mapping.jsonl +├── pesmaker_interval.yaml +└── toolkit.log +``` + +等间距采样不计算 FPS 描述符，因此不会生成 `selection_features.npy` 或 `fps_selection.png`。 + +## WSL 调用示例 + +Windows D 盘在 WSL 中通常映射到 `/mnt/d`。例如： + +```bash +conda run -n pesmaker python \ + /mnt/d/ResearchData/PESMaker_litupVer/PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py \ + --aimd-dir /mnt/d/ResearchData/example_AIMD \ + --interval 100 \ + --start-frame 0 +``` + +查看命令行帮助： + +```bash +python aimd_fps_to_nep.py --help +python aimd_interval_to_nep.py --help +``` + +## 输出文件说明 + +| 文件 | 用途 | +| --- | --- | +| `train.xyz` | 最终 NEP 训练数据，每帧包含总能量、力和总 Virial | +| `selected.xyz` | PESMaker 选中的无标签构型 | +| `manifest.jsonl` | PESMaker 选择清单，包含原始零基 `source_frame` | +| `frame_mapping.jsonl` | `train.xyz` 帧与 XDATCAR/OUTCAR 帧号的对应关系 | +| `pesmaker_fps.yaml` | FPS 脚本自动生成的 PESMaker 配置 | +| `pesmaker_interval.yaml` | 等间距脚本自动生成的 PESMaker 配置 | +| `selection_features.npy` | FPS 使用的简单几何特征 | +| `fps_selection.png` | FPS 选择过程图 | +| `toolkit.log` | 参数、路径、NBLOCK、帧数、标签范围、几何偏差和错误原因 | + +两个脚本都会拒绝覆盖已存在的身份目录，也不会自动创建编号子目录。 + +## 构型与标签检查 + +每个选中帧必须满足： + +- 原子数完全一致。 +- 元素及原子顺序完全一致。 +- 固定晶胞矩阵元素最大差不超过 `1e-4 Å`。 +- 周期性最小镜像下最大原子位移不超过 `1e-4 Å`。 +- `free_energy`、力和 stress 均存在且为有限值。 +- 力数组形状严格为 `N×3`。 +- 对应离子步能在 OUTCAR 中找到。 + +任一选中帧不满足要求时，整个转换失败，不会生成正式 `train.xyz`。 + +## train.xyz 标签定义 + +每个构型使用 GPUMD/NEP 扩展 XYZ 格式： + +```text +Lattice="ax ay az bx by bz cx cy cz" +Energy= +Properties=species:S:1:pos:R:3:force:R:3 +Virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" +pbc="T T T" +``` + +- `Energy`：VASP `free energy TOTEN`，单位 eV，是总能量。 +- `force`：每个原子的力，单位 eV/Å。 +- `Virial`：整个构型的 3×3 总 Virial，单位 eV，不是每原子值。 + +Virial 转换采用： + +```text +Virial = -V × stress_ASE +``` + +ASE 使用拉伸为正的 stress 约定，转换后的总 Virial 符合 GPUMD/NEP 使用的符号约定。 + +## 安全输出与失败行为 + +脚本先写入： + +```text +train.xyz.partial +frame_mapping.jsonl.partial +``` + +所有选中帧都匹配并通过标签检查后，才改名为正式文件： + +```text +train.xyz +frame_mapping.jsonl +``` + +如果失败： + +- 不生成正式 `train.xyz`。 +- 已完成的选择结果和 `toolkit.log` 会保留。 +- `.partial` 文件可能保留，但不能用于训练。 +- 失败原因会写入日志并显示在终端。 + +## 常见错误 + +### 输出目录已经存在 + +脚本不会覆盖已有结果： + +```text +Error: output directory already exists; refusing to overwrite +``` + +请先检查已有 `train.xyz` 和 `toolkit.log`，再决定是否保留或移走目录。 + +### 构型不匹配 + +常见原因包括： + +- XDATCAR 与 OUTCAR 不是同一次计算。 +- 文件被拼接或截断。 +- INCAR 中的 `NBLOCK` 与实际轨迹写出不一致。 +- AIMD 使用变胞设置。 +- 原子顺序被其他工具修改。 + +### 缺少能量、力或 stress + +脚本不会用其他能量定义代替 TOTEN，也不会生成缺少 Virial 的训练帧。应检查 OUTCAR 是否完整，以及相应离子步是否正常结束。 + +## 已完成的真实数据验证 + +FPS 脚本已在一个固定晶胞 VASP AIMD 目录上完成端到端测试： + +- XDATCAR：16,851 帧。 +- OUTCAR：约 1.46 GB。 +- FPS 测试选择：5 帧。 +- 选中离子步：1、5883、10596、14036、16583。 +- 每帧原子数：156。 +- `train.xyz` 可由 ASE 重新读取。 +- TOTEN 与 PESMaker `collect` 解析结果一致。 +- Virial 最大差约 `6.54×10⁻⁵ eV`，来自 OUTCAR 文本打印精度。 + +等间距脚本已在同一条 16,851 帧轨迹上使用以下参数完成端到端测试： + +```text +--start-frame 100 +--end-frame 16500 +--interval 4000 +``` + +验证结果： + +- 实际源帧：100、4100、8100、12100、16100。 +- 对应 OUTCAR 离子步：101、4101、8101、12101、16101（`NBLOCK=1`）。 +- `manifest.jsonl`、`frame_mapping.jsonl` 和预期帧号完全一致。 +- 5 个结构均包含有限的 Energy、`156×3` force 和 9 个 Virial 分量。 +- `train.xyz` 可由 ASE 重新读取。 +- 没有残留 `.partial` 文件。 + +真实测试只能证明对应数据及当前适用范围内的流程可用，不能替代对其他 VASP 版本、变胞或拼接轨迹的单独验证。 diff --git a/PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py b/PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py new file mode 100644 index 0000000..68117ef --- /dev/null +++ b/PESMaker_AIMD_Toolkit/aimd_interval_to_nep.py @@ -0,0 +1,386 @@ +#!/usr/bin/env python3 +"""Create a NEP training set by interval-sampling one VASP AIMD run. + +The script calls the public ``python -m pesmaker select`` command for +interval selection, then reuses the validated OUTCAR streaming and label +conversion helpers from ``aimd_fps_to_nep.py`` in this toolkit directory. +""" + +from __future__ import annotations + +import argparse +import logging +from pathlib import Path +import re +import subprocess +import sys + +from aimd_fps_to_nep import ( + ExtractionStats, + SelectionEntry, + ToolkitError, + extract_train_xyz, + load_selections, + parse_nblock, +) + + +OUTPUT_DIR_NAME = "PESMakerToolkit_AIMD_Interval_to_NEP" +REQUIRED_SELECTION_FILES = ("selected.xyz", "manifest.jsonl") +CONFIGURATION_PATTERN = re.compile( + r"^\s*(?:Direct|Cartesian)\s+configuration\s*=", + re.IGNORECASE, +) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Interval-sample a VASP AIMD XDATCAR and stream matching " + "energy/force/virial labels from OUTCAR into NEP train.xyz." + ) + ) + parser.add_argument( + "--aimd-dir", + type=Path, + required=True, + help="VASP AIMD directory containing INCAR, XDATCAR, and OUTCAR.", + ) + parser.add_argument( + "--interval", + type=int, + required=True, + help="Select every Nth XDATCAR frame (must be positive).", + ) + parser.add_argument( + "--start-frame", + type=int, + default=0, + help="First eligible zero-based XDATCAR frame. Default: 0.", + ) + parser.add_argument( + "--end-frame", + type=int, + default=None, + help=( + "Inclusive last eligible zero-based XDATCAR frame. " + "Default: the final frame." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + try: + _validate_cli_values( + interval=args.interval, + start_frame=args.start_frame, + end_frame=args.end_frame, + ) + run_toolkit( + aimd_dir=args.aimd_dir, + interval=args.interval, + start_frame=args.start_frame, + end_frame=args.end_frame, + ) + except ToolkitError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 2 + except KeyboardInterrupt: + print("Error: interrupted by user", file=sys.stderr) + return 130 + return 0 + + +def run_toolkit( + *, + aimd_dir: Path, + interval: int, + start_frame: int = 0, + end_frame: int | None = None, +) -> Path: + """Run interval selection and streamed OUTCAR label extraction.""" + _validate_cli_values( + interval=interval, + start_frame=start_frame, + end_frame=end_frame, + ) + aimd_dir = aimd_dir.expanduser().resolve() + incar, xdatcar, outcar = _validate_inputs(aimd_dir) + output_dir = aimd_dir / OUTPUT_DIR_NAME + if output_dir.exists(): + raise ToolkitError( + f"output directory already exists; refusing to overwrite: {output_dir}" + ) + + frame_count = count_xdatcar_frames(xdatcar) + expected_source_frames, resolved_end_frame = resolve_source_frames( + frame_count=frame_count, + interval=interval, + start_frame=start_frame, + end_frame=end_frame, + ) + + output_dir.mkdir(parents=False, exist_ok=False) + logger = _configure_logger(output_dir / "toolkit.log") + logger.info("PESMaker AIMD interval to NEP toolkit") + logger.info("AIMD directory : %s", aimd_dir) + logger.info("XDATCAR : %s", xdatcar) + logger.info("OUTCAR : %s", outcar) + logger.info("Output : %s", output_dir) + logger.info("XDATCAR frames : %d", frame_count) + logger.info("interval : %d frame(s)", interval) + logger.info("start_frame : %d", start_frame) + logger.info("end_frame : %d (inclusive)", resolved_end_frame) + logger.info("expected frames : %d", len(expected_source_frames)) + + try: + nblock = parse_nblock(incar) + logger.info("NBLOCK : %d", nblock) + config_path = output_dir / "pesmaker_interval.yaml" + _write_pesmaker_config( + config_path, + xdatcar=xdatcar, + output_dir=output_dir, + interval=interval, + start_frame=start_frame, + max_count=len(expected_source_frames), + ) + _run_pesmaker_select(config_path, aimd_dir=aimd_dir, logger=logger) + _verify_selection_outputs(output_dir) + + selections = load_selections(output_dir, nblock=nblock) + verify_selected_source_frames(selections, expected_source_frames) + logger.info("Interval selected: %d frame(s)", len(selections)) + stats = extract_train_xyz( + outcar=outcar, + selections=selections, + output_dir=output_dir, + logger=logger, + ) + _log_success(logger, stats, output_dir / "train.xyz") + return output_dir / "train.xyz" + except Exception as exc: + if isinstance(exc, ToolkitError): + logger.error("Toolkit failed: %s", exc) + raise + logger.exception("Toolkit failed with an unexpected error") + raise ToolkitError(str(exc)) from exc + + +def _validate_cli_values( + *, + interval: int, + start_frame: int, + end_frame: int | None, +) -> None: + if interval < 1: + raise ToolkitError("--interval must be a positive integer") + if start_frame < 0: + raise ToolkitError("--start-frame must be zero or a positive integer") + if end_frame is not None: + if end_frame < 0: + raise ToolkitError("--end-frame must be zero or a positive integer") + if end_frame < start_frame: + raise ToolkitError("--end-frame must not be smaller than --start-frame") + + +def _validate_inputs(aimd_dir: Path) -> tuple[Path, Path, Path]: + if not aimd_dir.is_dir(): + raise ToolkitError(f"AIMD directory does not exist: {aimd_dir}") + paths = tuple(aimd_dir / name for name in ("INCAR", "XDATCAR", "OUTCAR")) + missing = [path.name for path in paths if not path.is_file()] + if missing: + raise ToolkitError( + f"AIMD directory is missing required file(s): {', '.join(missing)}" + ) + return paths + + +def count_xdatcar_frames(xdatcar: Path) -> int: + """Count standard VASP XDATCAR configuration markers without loading frames.""" + frame_count = 0 + with xdatcar.open("r", encoding="utf-8", errors="replace") as handle: + for line in handle: + if CONFIGURATION_PATTERN.match(line): + frame_count += 1 + if frame_count < 1: + raise ToolkitError(f"XDATCAR contains no configuration frames: {xdatcar}") + return frame_count + + +def resolve_source_frames( + *, + frame_count: int, + interval: int, + start_frame: int, + end_frame: int | None, +) -> tuple[list[int], int]: + """Return the exact zero-based XDATCAR frames for inclusive bounds.""" + if frame_count < 1: + raise ToolkitError("XDATCAR frame count must be positive") + if start_frame >= frame_count: + raise ToolkitError( + f"--start-frame {start_frame} is outside XDATCAR frame range " + f"0..{frame_count - 1}" + ) + resolved_end = frame_count - 1 if end_frame is None else end_frame + if resolved_end >= frame_count: + raise ToolkitError( + f"--end-frame {resolved_end} is outside XDATCAR frame range " + f"0..{frame_count - 1}" + ) + source_frames = list(range(start_frame, resolved_end + 1, interval)) + if not source_frames: + raise ToolkitError("the requested interval range selected no XDATCAR frames") + return source_frames, resolved_end + + +def _configure_logger(path: Path) -> logging.Logger: + logger = logging.getLogger(f"pesmaker_aimd_interval_toolkit.{id(path)}") + logger.setLevel(logging.INFO) + logger.propagate = False + file_handler = logging.FileHandler(path, encoding="utf-8") + file_handler.setFormatter( + logging.Formatter("%(asctime)s %(levelname)s %(message)s") + ) + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(logging.Formatter("%(message)s")) + logger.addHandler(file_handler) + logger.addHandler(stream_handler) + return logger + + +def _write_pesmaker_config( + path: Path, + *, + xdatcar: Path, + output_dir: Path, + interval: int, + start_frame: int, + max_count: int, +) -> None: + try: + import yaml + except ImportError as exc: # pragma: no cover - PESMaker requires PyYAML + raise ToolkitError("PyYAML is required to generate the interval config") from exc + config = { + "project": "aimd_interval_to_nep", + "sampling": { + "engine": "none", + "selection": { + "method": "interval", + "trajectory_pattern": str(xdatcar), + "trajectory_format": "vasp-xdatcar", + "output_dir": str(output_dir), + "interval": interval, + "offset": start_frame, + "max_count": max_count, + "separate_trajectories": False, + }, + }, + } + path.write_text( + yaml.safe_dump(config, sort_keys=False, allow_unicode=True), + encoding="utf-8", + ) + + +def _run_pesmaker_select( + config_path: Path, + *, + aimd_dir: Path, + logger: logging.Logger, +) -> None: + command = [sys.executable, "-m", "pesmaker", "select", str(config_path)] + logger.info("Running : %s", " ".join(command)) + try: + process = subprocess.Popen( + command, + cwd=aimd_dir, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + ) + except OSError as exc: + raise ToolkitError(f"could not start PESMaker: {exc}") from exc + assert process.stdout is not None + try: + for line in process.stdout: + logger.info("[pesmaker] %s", line.rstrip()) + returncode = process.wait() + except KeyboardInterrupt: + process.terminate() + process.wait() + raise + if returncode != 0: + raise ToolkitError(f"PESMaker select failed with exit status {returncode}") + + +def _verify_selection_outputs(output_dir: Path) -> None: + missing = [ + name for name in REQUIRED_SELECTION_FILES if not (output_dir / name).is_file() + ] + if missing: + raise ToolkitError( + "PESMaker select completed without required output file(s): " + + ", ".join(missing) + ) + + +def verify_selected_source_frames( + selections: list[SelectionEntry], + expected_source_frames: list[int], +) -> None: + actual_source_frames = [selection.source_frame for selection in selections] + if actual_source_frames != expected_source_frames: + mismatch_index = next( + ( + index + for index, (actual, expected) in enumerate( + zip(actual_source_frames, expected_source_frames) + ) + if actual != expected + ), + min(len(actual_source_frames), len(expected_source_frames)), + ) + expected_value = ( + expected_source_frames[mismatch_index] + if mismatch_index < len(expected_source_frames) + else "" + ) + actual_value = ( + actual_source_frames[mismatch_index] + if mismatch_index < len(actual_source_frames) + else "" + ) + raise ToolkitError( + "PESMaker interval selection did not match the requested frame range: " + f"expected {len(expected_source_frames)} frame(s), got " + f"{len(actual_source_frames)}; first mismatch at selection " + f"{mismatch_index}: expected {expected_value}, got {actual_value}" + ) + + +def _log_success( + logger: logging.Logger, + stats: ExtractionStats, + train_path: Path, +) -> None: + logger.info("Dataset complete") + logger.info("Structures : %d", stats.written) + logger.info("Energy range : %.16g .. %.16g eV", stats.min_energy, stats.max_energy) + logger.info("Maximum force : %.16g eV/A", stats.max_force) + logger.info("Maximum |Virial| : %.16g eV", stats.max_abs_virial) + logger.info("Maximum cell diff: %.16g A", stats.max_cell_difference) + logger.info("Maximum pos diff : %.16g A", stats.max_position_difference) + logger.info("train.xyz : %s", train_path) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_aimd_interval_to_nep_toolkit.py b/tests/test_aimd_interval_to_nep_toolkit.py new file mode 100644 index 0000000..be8f2f9 --- /dev/null +++ b/tests/test_aimd_interval_to_nep_toolkit.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +import sys +from types import SimpleNamespace + +import pytest +import yaml + + +TOOLKIT_DIR = Path(__file__).resolve().parents[1] / "PESMaker_AIMD_Toolkit" +if str(TOOLKIT_DIR) not in sys.path: + sys.path.insert(0, str(TOOLKIT_DIR)) +TOOL_PATH = TOOLKIT_DIR / "aimd_interval_to_nep.py" +SPEC = importlib.util.spec_from_file_location("aimd_interval_to_nep", TOOL_PATH) +assert SPEC is not None and SPEC.loader is not None +tool = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = tool +SPEC.loader.exec_module(tool) + + +def test_count_xdatcar_frames_streams_configuration_markers(tmp_path: Path): + xdatcar = tmp_path / "XDATCAR" + xdatcar.write_text( + "header\nDirect configuration= 1\npositions\n" + " Cartesian configuration= 2\npositions\n", + encoding="utf-8", + ) + assert tool.count_xdatcar_frames(xdatcar) == 2 + + +def test_resolve_source_frames_uses_inclusive_end_without_forcing_it(): + frames, resolved_end = tool.resolve_source_frames( + frame_count=20, + interval=4, + start_frame=2, + end_frame=13, + ) + assert frames == [2, 6, 10] + assert resolved_end == 13 + + +def test_resolve_source_frames_defaults_to_final_frame(): + frames, resolved_end = tool.resolve_source_frames( + frame_count=11, + interval=5, + start_frame=0, + end_frame=None, + ) + assert frames == [0, 5, 10] + assert resolved_end == 10 + + +@pytest.mark.parametrize( + ("interval", "start_frame", "end_frame", "message"), + [ + (0, 0, None, "interval"), + (1, -1, None, "start-frame"), + (1, 5, 4, "end-frame"), + ], +) +def test_cli_value_validation(interval, start_frame, end_frame, message): + with pytest.raises(tool.ToolkitError, match=message): + tool._validate_cli_values( + interval=interval, + start_frame=start_frame, + end_frame=end_frame, + ) + + +def test_explicit_end_frame_outside_trajectory_is_rejected(): + with pytest.raises(tool.ToolkitError, match="outside XDATCAR frame range"): + tool.resolve_source_frames( + frame_count=10, + interval=2, + start_frame=0, + end_frame=10, + ) + + +def test_generated_interval_yaml_uses_exact_bounds(tmp_path: Path): + xdatcar = (tmp_path / "XDATCAR").resolve() + output_dir = (tmp_path / tool.OUTPUT_DIR_NAME).resolve() + output_dir.mkdir() + config_path = output_dir / "pesmaker_interval.yaml" + tool._write_pesmaker_config( + config_path, + xdatcar=xdatcar, + output_dir=output_dir, + interval=100, + start_frame=20, + max_count=5, + ) + + config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + selection = config["sampling"]["selection"] + assert selection == { + "method": "interval", + "trajectory_pattern": str(xdatcar), + "trajectory_format": "vasp-xdatcar", + "output_dir": str(output_dir), + "interval": 100, + "offset": 20, + "max_count": 5, + "separate_trajectories": False, + } + + +def test_manifest_source_frames_must_match_requested_sequence(): + selections = [ + SimpleNamespace(source_frame=2), + SimpleNamespace(source_frame=6), + SimpleNamespace(source_frame=10), + ] + tool.verify_selected_source_frames(selections, [2, 6, 10]) + with pytest.raises(tool.ToolkitError, match="did not match"): + tool.verify_selected_source_frames(selections, [2, 6, 11]) + + +def test_existing_output_directory_is_refused(tmp_path: Path): + for name in ("INCAR", "XDATCAR", "OUTCAR"): + (tmp_path / name).write_text("", encoding="utf-8") + (tmp_path / tool.OUTPUT_DIR_NAME).mkdir() + with pytest.raises(tool.ToolkitError, match="refusing to overwrite"): + tool.run_toolkit( + aimd_dir=tmp_path, + interval=10, + start_frame=0, + end_frame=None, + ) From 913c6a5ad6f8d8da992f94b4a78f704010ccfec4 Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Thu, 23 Jul 2026 14:28:35 +0800 Subject: [PATCH 7/8] Add AIMD train.xyz merge toolkit --- PESMaker_AIMD_Toolkit/README.md | 138 +++++ PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py | 509 ++++++++++++++++++ tests/test_merge_aimd_train_xyz_toolkit.py | 170 ++++++ 3 files changed, 817 insertions(+) create mode 100644 PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py create mode 100644 tests/test_merge_aimd_train_xyz_toolkit.py diff --git a/PESMaker_AIMD_Toolkit/README.md b/PESMaker_AIMD_Toolkit/README.md index 912aa24..65449de 100644 --- a/PESMaker_AIMD_Toolkit/README.md +++ b/PESMaker_AIMD_Toolkit/README.md @@ -6,6 +6,7 @@ | --- | --- | --- | | `aimd_fps_to_nep.py` | PESMaker 最远点采样（FPS） | 根据构型差异选择代表性结构 | | `aimd_interval_to_nep.py` | 等间距采样 | 按固定轨迹帧间隔抽稀 AIMD 时间序列 | +| `merge_aimd_train_xyz.py` | 指定来源合并 | 将多个正式 FPS 或 Interval `train.xyz` 合并并记录来源范围 | 两个脚本都不会修改或导入 PESMaker 私有源码，而是调用正式命令： @@ -214,6 +215,142 @@ PESMakerToolkit_AIMD_Interval_to_NEP/ 等间距采样不计算 FPS 描述符，因此不会生成 `selection_features.npy` 或 `fps_selection.png`。 +## 合并脚本用途与语法 + +`merge_aimd_train_xyz.py` 用于合并多个 AIMD 子目录中已经生成并验证的正式 `train.xyz`。它不重新解析 AIMD、不调用 PESMaker，也不修改任何 PESMaker 源码。 + +用户必须明确指定合并来源，不允许脚本自动猜测或混合 FPS 与 Interval： + +```bash +python PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py \ + --aimd-root /absolute/path/to/AIMD_ROOT \ + --source interval +``` + +合并 FPS 来源： + +```bash +python PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py \ + --aimd-root /absolute/path/to/AIMD_ROOT \ + --source fps +``` + +参数： + +| 参数 | 是否必填 | 说明 | +| --- | --- | --- | +| `--aimd-root` | 是 | 直接包含多个 VASP AIMD 运行目录的根目录 | +| `--source` | 是 | 只能是 `interval` 或 `fps`，不提供默认值 | + +### 预期输入结构 + +`--aimd-root` 的直接子目录中，正式 AIMD 目录必须同时包含 INCAR、XDATCAR 和 OUTCAR： + +```text +AIMD_ROOT/ +├── aimd_case_01/ +│ ├── INCAR +│ ├── XDATCAR +│ ├── OUTCAR +│ ├── PESMakerToolkit_AIMD_Interval_to_NEP/ +│ │ └── train.xyz +│ └── PESMakerToolkit_AIMD_FPS_to_NEP/ +│ └── train.xyz +├── aimd_case_02/ +│ ├── INCAR +│ ├── XDATCAR +│ ├── OUTCAR +│ └── ... +└── scripts/ # 不含三项 VASP 文件，因此忽略 +``` + +来源参数和实际查找路径严格对应： + +```text +--source interval +→ /PESMakerToolkit_AIMD_Interval_to_NEP/train.xyz + +--source fps +→ /PESMakerToolkit_AIMD_FPS_to_NEP/train.xyz +``` + +脚本只检查 AIMD 根目录的直接子目录，不递归搜索任意 `train.xyz`。因此不会误收集 `_smoke_test`、`_failed_*` 或其他临时目录。 + +每一个识别出的正式 AIMD 子目录都必须存在所选来源的 `train.xyz`。如果任一目录缺失，脚本会在创建合并输出前失败并列出缺失目录，不会静默跳过。 + +### 合并输出目录 + +输出位于 AIMD 根目录的同级目录，名称明确包含来源： + +```text +/Merged_Interval_NEP_TrainXYZ/ +/Merged_FPS_NEP_TrainXYZ/ +``` + +例如： + +```text +/mnt/d/ResearchData/ +├── growthDatasetPreparation_AIMD/ +└── Merged_Interval_NEP_TrainXYZ/ +``` + +输出文件结构： + +```text +Merged_Interval_NEP_TrainXYZ/ +├── train.xyz +├── source_ranges.tsv +├── README.md +└── merge.log +``` + +- `train.xyz`：按照 AIMD 目录名称不区分大小写排序后直接拼接的训练集，不重新格式化标签。 +- `source_ranges.tsv`：机器可读的来源、帧数、元素、原子数和合并范围。 +- `README.md`：自动生成的人类可读来源表、目录说明、合并规则和复现命令。 +- `merge.log`：来源验证、帧数范围、总帧数和失败原因。 + +`source_ranges.tsv` 包含： + +```text +order +sampling_source +aimd_directory +source_train_xyz +frame_count +atom_count +elements +merged_start_0based +merged_end_0based +merged_start_1based +merged_end_1based +``` + +零基和一基范围均为包含式。脚本根据实际读取帧数动态计算范围，不硬编码当前数据集的数量。 + +### 合并验证与安全行为 + +每个来源和最终合并文件均使用 ASE 流式验证： + +- extxyz 至少包含一帧。 +- 晶胞、位置、Energy、force 和 Virial 均存在且为有限值。 +- force 形状严格为 `N×3`。 +- Virial 恰好包含 9 个分量。 +- 合并总帧数等于所有来源帧数之和。 +- 合并元素集合与来源元素集合一致。 + +合并过程先生成： + +```text +train.xyz.partial +source_ranges.tsv.partial +README.md.partial +``` + +全部验证通过后才改名为正式文件，其中 `train.xyz` 最后改名。输出目录已经存在时脚本拒绝覆盖。 + +第一版不会执行随机打乱、去重、二次 FPS、train/test 划分或标签单位转换。 + ## WSL 调用示例 Windows D 盘在 WSL 中通常映射到 `/mnt/d`。例如： @@ -231,6 +368,7 @@ conda run -n pesmaker python \ ```bash python aimd_fps_to_nep.py --help python aimd_interval_to_nep.py --help +python merge_aimd_train_xyz.py --help ``` ## 输出文件说明 diff --git a/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py b/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py new file mode 100644 index 0000000..1f2d4b0 --- /dev/null +++ b/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py @@ -0,0 +1,509 @@ +#!/usr/bin/env python3 +"""Merge validated AIMD toolkit train.xyz files from one sampling source. + +The user must choose ``interval`` or ``fps`` explicitly. The script scans only +direct AIMD children and only the official output directory for that source; +it never searches recursively for arbitrary files named ``train.xyz``. +""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import dataclass +from datetime import datetime +import logging +import math +from pathlib import Path +import sys +from typing import Any, Iterable + +import numpy as np + + +@dataclass(frozen=True) +class SourceMode: + """Directory mapping and display name for one sampling source.""" + + key: str + display_name: str + output_directory: str + + +SOURCE_MODES = { + "interval": SourceMode( + key="interval", + display_name="Interval", + output_directory="PESMakerToolkit_AIMD_Interval_to_NEP", + ), + "fps": SourceMode( + key="fps", + display_name="FPS", + output_directory="PESMakerToolkit_AIMD_FPS_to_NEP", + ), +} + + +@dataclass(frozen=True) +class DatasetStats: + """Labels and structure information observed in one extxyz dataset.""" + + frame_count: int + min_atom_count: int + max_atom_count: int + elements: tuple[str, ...] + + +@dataclass(frozen=True) +class SourceSummary: + """One source dataset and its inclusive range in the merged dataset.""" + + order: int + aimd_directory: Path + train_xyz: Path + stats: DatasetStats + merged_start_0based: int + merged_end_0based: int + + +class MergeError(RuntimeError): + """A user-facing merge failure.""" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Merge validated train.xyz files from the explicitly selected " + "AIMD toolkit sampling source." + ) + ) + parser.add_argument( + "--aimd-root", + type=Path, + required=True, + help="Directory whose direct children are VASP AIMD run directories.", + ) + parser.add_argument( + "--source", + choices=tuple(SOURCE_MODES), + required=True, + help="Sampling source to merge: interval or fps.", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + try: + run_merge(aimd_root=args.aimd_root, source=args.source) + except MergeError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 2 + except KeyboardInterrupt: + print("Error: interrupted by user", file=sys.stderr) + return 130 + return 0 + + +def run_merge(*, aimd_root: Path, source: str) -> Path: + """Validate and merge all direct AIMD children for one sampling source.""" + if source not in SOURCE_MODES: + raise MergeError( + f"unsupported source {source!r}; choose one of: " + + ", ".join(SOURCE_MODES) + ) + mode = SOURCE_MODES[source] + aimd_root = aimd_root.expanduser().resolve() + if not aimd_root.is_dir(): + raise MergeError(f"AIMD root directory does not exist: {aimd_root}") + + output_dir = aimd_root.parent / f"Merged_{mode.display_name}_NEP_TrainXYZ" + if output_dir.exists(): + raise MergeError( + f"output directory already exists; refusing to overwrite: {output_dir}" + ) + + aimd_directories = discover_aimd_directories(aimd_root) + source_paths = require_source_train_xyz(aimd_directories, mode) + + output_dir.mkdir(parents=False, exist_ok=False) + logger = configure_logger(output_dir / "merge.log") + logger.info("PESMaker AIMD train.xyz merge toolkit") + logger.info("AIMD root : %s", aimd_root) + logger.info("Sampling source : %s", mode.key) + logger.info("Source directory: %s", mode.output_directory) + logger.info("AIMD directories: %d", len(aimd_directories)) + logger.info("Output : %s", output_dir) + + train_partial = output_dir / "train.xyz.partial" + ranges_partial = output_dir / "source_ranges.tsv.partial" + readme_partial = output_dir / "README.md.partial" + + try: + summaries = build_source_summaries( + aimd_root=aimd_root, + source_paths=source_paths, + logger=logger, + ) + total_frames = sum(summary.stats.frame_count for summary in summaries) + concatenate_train_xyz( + [summary.train_xyz for summary in summaries], + train_partial, + ) + write_source_ranges( + ranges_partial, + aimd_root=aimd_root, + mode=mode, + summaries=summaries, + ) + write_output_readme( + readme_partial, + aimd_root=aimd_root, + mode=mode, + summaries=summaries, + total_frames=total_frames, + ) + + merged_stats = validate_train_xyz(train_partial) + if merged_stats.frame_count != total_frames: + raise MergeError( + "merged frame-count mismatch: expected " + f"{total_frames}, got {merged_stats.frame_count}" + ) + expected_elements = tuple( + sorted( + { + element + for summary in summaries + for element in summary.stats.elements + } + ) + ) + if merged_stats.elements != expected_elements: + raise MergeError( + "merged element set does not match the source datasets: " + f"expected {expected_elements}, got {merged_stats.elements}" + ) + + ranges_partial.replace(output_dir / "source_ranges.tsv") + readme_partial.replace(output_dir / "README.md") + train_partial.replace(output_dir / "train.xyz") + logger.info("Merge complete") + logger.info("Source datasets : %d", len(summaries)) + logger.info("Total frames : %d", total_frames) + logger.info("Elements : %s", " ".join(merged_stats.elements)) + logger.info("train.xyz : %s", output_dir / "train.xyz") + return output_dir / "train.xyz" + except Exception as exc: + if isinstance(exc, MergeError): + logger.error("Merge failed: %s", exc) + raise + logger.exception("Merge failed with an unexpected error") + raise MergeError(str(exc)) from exc + + +def discover_aimd_directories(aimd_root: Path) -> list[Path]: + """Find direct children that contain the three required VASP files.""" + required_files = ("INCAR", "XDATCAR", "OUTCAR") + directories = [ + child + for child in aimd_root.iterdir() + if child.is_dir() + and all((child / filename).is_file() for filename in required_files) + ] + directories.sort(key=lambda path: path.name.casefold()) + if not directories: + raise MergeError( + "AIMD root has no direct child containing INCAR, XDATCAR, and OUTCAR: " + f"{aimd_root}" + ) + return directories + + +def require_source_train_xyz( + aimd_directories: list[Path], + mode: SourceMode, +) -> list[Path]: + """Require every discovered AIMD child to have the chosen official source.""" + source_paths = [ + directory / mode.output_directory / "train.xyz" + for directory in aimd_directories + ] + missing = [ + directory + for directory, source_path in zip(aimd_directories, source_paths) + if not source_path.is_file() + ] + if missing: + names = ", ".join(directory.name for directory in missing) + raise MergeError( + f"{len(missing)} AIMD directorie(s) lack the selected " + f"{mode.key} train.xyz in {mode.output_directory}: {names}" + ) + return source_paths + + +def build_source_summaries( + *, + aimd_root: Path, + source_paths: list[Path], + logger: logging.Logger, +) -> list[SourceSummary]: + """Validate source datasets and assign contiguous merged frame ranges.""" + summaries = [] + next_start = 0 + for order, train_xyz in enumerate(source_paths, start=1): + stats = validate_train_xyz(train_xyz) + merged_end = next_start + stats.frame_count - 1 + summary = SourceSummary( + order=order, + aimd_directory=train_xyz.parent.parent, + train_xyz=train_xyz, + stats=stats, + merged_start_0based=next_start, + merged_end_0based=merged_end, + ) + summaries.append(summary) + logger.info( + "Source %d : %s | %d frame(s) | merged %d..%d", + order, + train_xyz.relative_to(aimd_root).as_posix(), + stats.frame_count, + next_start, + merged_end, + ) + next_start = merged_end + 1 + return summaries + + +def validate_train_xyz(path: Path) -> DatasetStats: + """Stream one extxyz dataset and require finite NEP labels on every frame.""" + try: + from ase.io import iread + except ImportError as exc: # pragma: no cover - toolkit requires ASE + raise MergeError("ASE is required to validate train.xyz") from exc + + frame_count = 0 + min_atom_count = sys.maxsize + max_atom_count = 0 + elements: set[str] = set() + try: + frames: Iterable[Any] = iread(path, index=":", format="extxyz") + for frame_index, atoms in enumerate(frames): + validate_nep_frame(atoms, path=path, frame_index=frame_index) + atom_count = len(atoms) + frame_count += 1 + min_atom_count = min(min_atom_count, atom_count) + max_atom_count = max(max_atom_count, atom_count) + elements.update(atoms.get_chemical_symbols()) + except MergeError: + raise + except Exception as exc: + raise MergeError(f"could not read extxyz dataset {path}: {exc}") from exc + + if frame_count == 0: + raise MergeError(f"train.xyz contains no frames: {path}") + return DatasetStats( + frame_count=frame_count, + min_atom_count=min_atom_count, + max_atom_count=max_atom_count, + elements=tuple(sorted(elements)), + ) + + +def validate_nep_frame(atoms: Any, *, path: Path, frame_index: int) -> None: + """Validate labels and geometry required by GPUMD/NEP extxyz input.""" + prefix = f"{path} frame {frame_index}" + if len(atoms) < 1: + raise MergeError(f"{prefix} has no atoms") + cell = np.asarray(atoms.cell.array, dtype=float) + positions = np.asarray(atoms.positions, dtype=float) + if cell.shape != (3, 3) or not np.all(np.isfinite(cell)): + raise MergeError(f"{prefix} has a non-finite or invalid cell") + if positions.shape != (len(atoms), 3) or not np.all(np.isfinite(positions)): + raise MergeError(f"{prefix} has non-finite or invalid positions") + + if "Energy" not in atoms.info: + raise MergeError(f"{prefix} has no Energy label") + energy = float(atoms.info["Energy"]) + if not math.isfinite(energy): + raise MergeError(f"{prefix} has a non-finite Energy label") + + if "force" not in atoms.arrays: + raise MergeError(f"{prefix} has no force labels") + forces = np.asarray(atoms.arrays["force"], dtype=float) + if forces.shape != (len(atoms), 3) or not np.all(np.isfinite(forces)): + raise MergeError(f"{prefix} has non-finite or invalid force labels") + + if "Virial" not in atoms.info: + raise MergeError(f"{prefix} has no Virial label") + virial = np.asarray(atoms.info["Virial"], dtype=float) + if virial.size != 9 or not np.all(np.isfinite(virial)): + raise MergeError(f"{prefix} has non-finite or invalid Virial labels") + + +def concatenate_train_xyz(source_paths: list[Path], destination: Path) -> None: + """Concatenate validated source text without reformatting extxyz frames.""" + with destination.open("x", encoding="utf-8", newline="\n") as output: + for source_path in source_paths: + last_character = "\n" + with source_path.open("r", encoding="utf-8") as source: + while True: + chunk = source.read(1024 * 1024) + if not chunk: + break + output.write(chunk) + last_character = chunk[-1] + if last_character != "\n": + output.write("\n") + + +def write_source_ranges( + path: Path, + *, + aimd_root: Path, + mode: SourceMode, + summaries: list[SourceSummary], +) -> None: + """Write a machine-readable source and merged-range table.""" + with path.open("x", encoding="utf-8", newline="") as handle: + writer = csv.writer(handle, delimiter="\t", lineterminator="\n") + writer.writerow( + ( + "order", + "sampling_source", + "aimd_directory", + "source_train_xyz", + "frame_count", + "atom_count", + "elements", + "merged_start_0based", + "merged_end_0based", + "merged_start_1based", + "merged_end_1based", + ) + ) + for summary in summaries: + writer.writerow( + ( + summary.order, + mode.key, + summary.aimd_directory.name, + summary.train_xyz.relative_to(aimd_root).as_posix(), + summary.stats.frame_count, + atom_count_label(summary.stats), + " ".join(summary.stats.elements), + summary.merged_start_0based, + summary.merged_end_0based, + summary.merged_start_0based + 1, + summary.merged_end_0based + 1, + ) + ) + + +def write_output_readme( + path: Path, + *, + aimd_root: Path, + mode: SourceMode, + summaries: list[SourceSummary], + total_frames: int, +) -> None: + """Write a concise human-readable provenance and layout description.""" + timestamp = datetime.now().astimezone().isoformat(timespec="seconds") + lines = [ + f"# Merged {mode.display_name} NEP train.xyz", + "", + "本目录由 `merge_aimd_train_xyz.py` 自动生成。", + "", + f"- AIMD 根目录：`{aimd_root}`", + f"- 采样来源：`{mode.key}`", + f"- 来源输出目录：`{mode.output_directory}`", + f"- 合并时间：`{timestamp}`", + f"- 来源轨迹数：`{len(summaries)}`", + f"- 合并总帧数：`{total_frames}`", + "", + "## 文件结构", + "", + "```text", + f"Merged_{mode.display_name}_NEP_TrainXYZ/", + "├── train.xyz", + "├── source_ranges.tsv", + "├── README.md", + "└── merge.log", + "```", + "", + "- `train.xyz`：按下表顺序拼接的 GPUMD/NEP 训练集。", + "- `source_ranges.tsv`：机器可读的来源路径、帧数和合并范围。", + "- `README.md`：本说明。", + "- `merge.log`：验证、合并过程和错误日志。", + "", + "## 来源与帧范围", + "", + "范围均为包含式。零基范围适合程序处理，一基范围适合人工计数。", + "", + "| 顺序 | AIMD 目录 | 帧数 | 原子数 | 元素 | 零基范围 | 一基范围 |", + "| ---: | --- | ---: | ---: | --- | ---: | ---: |", + ] + for summary in summaries: + lines.append( + "| " + f"{summary.order} | {markdown_cell(summary.aimd_directory.name)} | " + f"{summary.stats.frame_count} | {atom_count_label(summary.stats)} | " + f"{' '.join(summary.stats.elements)} | " + f"{summary.merged_start_0based}–{summary.merged_end_0based} | " + f"{summary.merged_start_0based + 1}–{summary.merged_end_0based + 1} |" + ) + lines.extend( + [ + "", + "## 合并规则", + "", + "- 只读取 AIMD 根目录的直接子目录。", + f"- 只读取 `{mode.output_directory}/train.xyz`。", + "- 不递归搜索其他 `train.xyz`，不混合其他采样来源。", + "- 来源按 AIMD 目录名称不区分大小写排序。", + "- 不随机打乱、不去重、不划分 train/test、不转换标签单位。", + "- 来源文件只读，合并过程不会修改任何原始数据。", + "", + "## 复现语法", + "", + "```bash", + "python merge_aimd_train_xyz.py \\", + f" --aimd-root {aimd_root} \\", + f" --source {mode.key}", + "```", + "", + ] + ) + path.write_text("\n".join(lines), encoding="utf-8") + + +def atom_count_label(stats: DatasetStats) -> str: + if stats.min_atom_count == stats.max_atom_count: + return str(stats.min_atom_count) + return f"{stats.min_atom_count}-{stats.max_atom_count}" + + +def markdown_cell(value: str) -> str: + return value.replace("|", "\\|") + + +def configure_logger(path: Path) -> logging.Logger: + logger = logging.getLogger(f"pesmaker_aimd_merge_toolkit.{id(path)}") + logger.setLevel(logging.INFO) + logger.propagate = False + file_handler = logging.FileHandler(path, encoding="utf-8") + file_handler.setFormatter( + logging.Formatter("%(asctime)s %(levelname)s %(message)s") + ) + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(logging.Formatter("%(message)s")) + logger.addHandler(file_handler) + logger.addHandler(stream_handler) + return logger + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_merge_aimd_train_xyz_toolkit.py b/tests/test_merge_aimd_train_xyz_toolkit.py new file mode 100644 index 0000000..256d42d --- /dev/null +++ b/tests/test_merge_aimd_train_xyz_toolkit.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import csv +import importlib.util +from pathlib import Path +import sys + +import numpy as np +import pytest +from ase import Atoms +from ase.io import read + + +TOOLKIT_DIR = Path(__file__).resolve().parents[1] / "PESMaker_AIMD_Toolkit" +if str(TOOLKIT_DIR) not in sys.path: + sys.path.insert(0, str(TOOLKIT_DIR)) + +MERGE_PATH = TOOLKIT_DIR / "merge_aimd_train_xyz.py" +MERGE_SPEC = importlib.util.spec_from_file_location( + "merge_aimd_train_xyz", MERGE_PATH +) +assert MERGE_SPEC is not None and MERGE_SPEC.loader is not None +merge_tool = importlib.util.module_from_spec(MERGE_SPEC) +sys.modules[MERGE_SPEC.name] = merge_tool +MERGE_SPEC.loader.exec_module(merge_tool) + +FPS_PATH = TOOLKIT_DIR / "aimd_fps_to_nep.py" +FPS_SPEC = importlib.util.spec_from_file_location("merge_test_fps_tool", FPS_PATH) +assert FPS_SPEC is not None and FPS_SPEC.loader is not None +fps_tool = importlib.util.module_from_spec(FPS_SPEC) +sys.modules[FPS_SPEC.name] = fps_tool +FPS_SPEC.loader.exec_module(fps_tool) + + +def make_aimd_directory(root: Path, name: str) -> Path: + directory = root / name + directory.mkdir() + for filename in ("INCAR", "XDATCAR", "OUTCAR"): + (directory / filename).write_text("test\n", encoding="utf-8") + return directory + + +def write_train_xyz(path: Path, energies: list[float], *, atom_count: int = 2): + path.parent.mkdir() + atoms = Atoms( + "H" * atom_count, + positions=np.arange(atom_count * 3, dtype=float).reshape(atom_count, 3), + cell=np.eye(3) * 10.0, + pbc=True, + ) + with path.open("w", encoding="utf-8") as handle: + for energy in energies: + fps_tool._write_nep_frame( + handle, + atoms=atoms, + energy=energy, + forces=np.zeros((atom_count, 3)), + virial=np.eye(3) * energy, + ) + + +def test_source_mode_maps_to_exact_official_directories(): + assert ( + merge_tool.SOURCE_MODES["interval"].output_directory + == "PESMakerToolkit_AIMD_Interval_to_NEP" + ) + assert ( + merge_tool.SOURCE_MODES["fps"].output_directory + == "PESMakerToolkit_AIMD_FPS_to_NEP" + ) + + +def test_discovery_uses_only_direct_vasp_aimd_children(tmp_path: Path): + root = tmp_path / "aimd" + root.mkdir() + valid_b = make_aimd_directory(root, "b_case") + valid_a = make_aimd_directory(root, "A_case") + (root / "scripts").mkdir() + nested = root / "container" / "nested" + nested.mkdir(parents=True) + for filename in ("INCAR", "XDATCAR", "OUTCAR"): + (nested / filename).write_text("test\n", encoding="utf-8") + + assert merge_tool.discover_aimd_directories(root) == [valid_a, valid_b] + + +def test_missing_selected_source_is_rejected_without_output(tmp_path: Path): + root = tmp_path / "aimd" + root.mkdir() + directory = make_aimd_directory(root, "case") + write_train_xyz( + directory + / "PESMakerToolkit_AIMD_Interval_to_NEP" + / "train.xyz", + [-1.0], + ) + + with pytest.raises(merge_tool.MergeError, match="lack the selected fps"): + merge_tool.run_merge(aimd_root=root, source="fps") + assert not (tmp_path / "Merged_FPS_NEP_TrainXYZ").exists() + + +def test_merge_interval_sources_end_to_end(tmp_path: Path): + root = tmp_path / "aimd" + root.mkdir() + case_b = make_aimd_directory(root, "b_case") + case_a = make_aimd_directory(root, "A_case") + source_directory = "PESMakerToolkit_AIMD_Interval_to_NEP" + write_train_xyz(case_a / source_directory / "train.xyz", [-1.0, -2.0]) + write_train_xyz(case_b / source_directory / "train.xyz", [-3.0], atom_count=3) + + train_path = merge_tool.run_merge(aimd_root=root, source="interval") + output = tmp_path / "Merged_Interval_NEP_TrainXYZ" + + assert train_path == output / "train.xyz" + assert {path.name for path in output.iterdir()} == { + "train.xyz", + "source_ranges.tsv", + "README.md", + "merge.log", + } + frames = read(train_path, index=":", format="extxyz") + assert [frame.info["Energy"] for frame in frames] == [-1.0, -2.0, -3.0] + assert [len(frame) for frame in frames] == [2, 2, 3] + + with (output / "source_ranges.tsv").open(encoding="utf-8", newline="") as handle: + records = list(csv.DictReader(handle, delimiter="\t")) + assert [record["aimd_directory"] for record in records] == ["A_case", "b_case"] + assert [record["frame_count"] for record in records] == ["2", "1"] + assert [record["merged_start_0based"] for record in records] == ["0", "2"] + assert [record["merged_end_0based"] for record in records] == ["1", "2"] + readme = (output / "README.md").read_text(encoding="utf-8") + assert "采样来源：`interval`" in readme + assert "0–1" in readme + assert "3" in readme + + +def test_invalid_source_dataset_does_not_create_formal_train(tmp_path: Path): + root = tmp_path / "aimd" + root.mkdir() + directory = make_aimd_directory(root, "case") + train_xyz = ( + directory + / "PESMakerToolkit_AIMD_Interval_to_NEP" + / "train.xyz" + ) + train_xyz.parent.mkdir() + train_xyz.write_text( + "1\n" + 'Lattice="10 0 0 0 10 0 0 0 10" Energy=-1 ' + 'Properties=species:S:1:pos:R:3:force:R:3 pbc="T T T"\n' + "H 0 0 0 0 0 0\n", + encoding="utf-8", + ) + + with pytest.raises(merge_tool.MergeError, match="no Virial label"): + merge_tool.run_merge(aimd_root=root, source="interval") + output = tmp_path / "Merged_Interval_NEP_TrainXYZ" + assert output.is_dir() + assert not (output / "train.xyz").exists() + assert (output / "merge.log").is_file() + + +def test_existing_output_directory_is_refused(tmp_path: Path): + root = tmp_path / "aimd" + root.mkdir() + make_aimd_directory(root, "case") + (tmp_path / "Merged_Interval_NEP_TrainXYZ").mkdir() + with pytest.raises(merge_tool.MergeError, match="refusing to overwrite"): + merge_tool.run_merge(aimd_root=root, source="interval") From c9442d3dda7844589d0dbc3b520ebe72503bc551 Mon Sep 17 00:00:00 2001 From: LETIANLI Date: Thu, 23 Jul 2026 14:44:04 +0800 Subject: [PATCH 8/8] Place merged datasets inside AIMD root --- PESMaker_AIMD_Toolkit/README.md | 11 ++++++----- PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py | 7 +++++-- tests/test_merge_aimd_train_xyz_toolkit.py | 8 ++++---- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/PESMaker_AIMD_Toolkit/README.md b/PESMaker_AIMD_Toolkit/README.md index 65449de..ce26db0 100644 --- a/PESMaker_AIMD_Toolkit/README.md +++ b/PESMaker_AIMD_Toolkit/README.md @@ -280,18 +280,19 @@ AIMD_ROOT/ ### 合并输出目录 -输出位于 AIMD 根目录的同级目录，名称明确包含来源： +输出位于 AIMD 根目录内部，与其中各个 AIMD 项目目录同级，名称明确包含来源： ```text -/Merged_Interval_NEP_TrainXYZ/ -/Merged_FPS_NEP_TrainXYZ/ +/Merged_Interval_NEP_TrainXYZ/ +/Merged_FPS_NEP_TrainXYZ/ ``` 例如： ```text -/mnt/d/ResearchData/ -├── growthDatasetPreparation_AIMD/ +/mnt/d/ResearchData/growthDatasetPreparation_AIMD/ +├── AIMD项目1/ +├── AIMD项目2/ └── Merged_Interval_NEP_TrainXYZ/ ``` diff --git a/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py b/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py index 1f2d4b0..20e7869 100644 --- a/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py +++ b/PESMaker_AIMD_Toolkit/merge_aimd_train_xyz.py @@ -81,7 +81,10 @@ def build_parser() -> argparse.ArgumentParser: "--aimd-root", type=Path, required=True, - help="Directory whose direct children are VASP AIMD run directories.", + help=( + "Directory whose direct children are VASP AIMD run directories; " + "the merged output directory is created inside this directory." + ), ) parser.add_argument( "--source", @@ -118,7 +121,7 @@ def run_merge(*, aimd_root: Path, source: str) -> Path: if not aimd_root.is_dir(): raise MergeError(f"AIMD root directory does not exist: {aimd_root}") - output_dir = aimd_root.parent / f"Merged_{mode.display_name}_NEP_TrainXYZ" + output_dir = aimd_root / f"Merged_{mode.display_name}_NEP_TrainXYZ" if output_dir.exists(): raise MergeError( f"output directory already exists; refusing to overwrite: {output_dir}" diff --git a/tests/test_merge_aimd_train_xyz_toolkit.py b/tests/test_merge_aimd_train_xyz_toolkit.py index 256d42d..5bf8dd7 100644 --- a/tests/test_merge_aimd_train_xyz_toolkit.py +++ b/tests/test_merge_aimd_train_xyz_toolkit.py @@ -97,7 +97,7 @@ def test_missing_selected_source_is_rejected_without_output(tmp_path: Path): with pytest.raises(merge_tool.MergeError, match="lack the selected fps"): merge_tool.run_merge(aimd_root=root, source="fps") - assert not (tmp_path / "Merged_FPS_NEP_TrainXYZ").exists() + assert not (root / "Merged_FPS_NEP_TrainXYZ").exists() def test_merge_interval_sources_end_to_end(tmp_path: Path): @@ -110,7 +110,7 @@ def test_merge_interval_sources_end_to_end(tmp_path: Path): write_train_xyz(case_b / source_directory / "train.xyz", [-3.0], atom_count=3) train_path = merge_tool.run_merge(aimd_root=root, source="interval") - output = tmp_path / "Merged_Interval_NEP_TrainXYZ" + output = root / "Merged_Interval_NEP_TrainXYZ" assert train_path == output / "train.xyz" assert {path.name for path in output.iterdir()} == { @@ -155,7 +155,7 @@ def test_invalid_source_dataset_does_not_create_formal_train(tmp_path: Path): with pytest.raises(merge_tool.MergeError, match="no Virial label"): merge_tool.run_merge(aimd_root=root, source="interval") - output = tmp_path / "Merged_Interval_NEP_TrainXYZ" + output = root / "Merged_Interval_NEP_TrainXYZ" assert output.is_dir() assert not (output / "train.xyz").exists() assert (output / "merge.log").is_file() @@ -165,6 +165,6 @@ def test_existing_output_directory_is_refused(tmp_path: Path): root = tmp_path / "aimd" root.mkdir() make_aimd_directory(root, "case") - (tmp_path / "Merged_Interval_NEP_TrainXYZ").mkdir() + (root / "Merged_Interval_NEP_TrainXYZ").mkdir() with pytest.raises(merge_tool.MergeError, match="refusing to overwrite"): merge_tool.run_merge(aimd_root=root, source="interval")