Add export_valid_split: write the valid_ratio split as GPUMD-ready files

README.md

@@ -180,6 +180,12 @@ train_nep("nep.in", "train.xyz", output_dir="output")
 # with a validation set (either form):
 train_nep("nep.in", "train.xyz", output_dir="output", valid_file="valid.xyz")
 train_nep("nep.in", "train.xyz", output_dir="output", valid_ratio=0.1)
+
+# export the exact valid_ratio split as GPUMD-ready files — train the same
+# partition in GPUMD (or anything else) and compare loss curves directly:
+from torchnep import export_valid_split
+export_valid_split("train.xyz", valid_ratio=0.1, run_seed=42,
+                   output_dir="split")   # writes split/train.xyz + split/test.xyz
 ```
 
 ```bash

releaseNotes.md

@@ -1,5 +1,12 @@
 # Release Notes
 
+## Unreleased
+
+- **`export_valid_split`**: write the exact `valid_ratio` split
+  `train_nep` uses as verbatim GPUMD-ready `train.xyz` / `test.xyz`
+  files, so the same data partition can be trained in GPUMD and the loss
+  curves compared directly.
+
 ## 1.0.2a1
 
 - **Streaming-only data path**: the preloaded GPU data store and the

tests/README.md

@@ -16,7 +16,7 @@ TEST_DEVICE=cpu pytest tests/      # restrict device (default: cpu + cuda if pre
 | `test_parsing.py` | Legacy and current `l_max` nep.in / nep.txt parsing. |
 | `test_ase_calculator.py` | Optional ASE calculator (energy/forces/stress, ZBL split). |
 | `test_b1_and_gpumd_qscaler.py` | Analytical `b1` offset (residual → 0, `nep_best` ≤ `nep_final`); `use_gpumd_qscaler` reproduces GPUMD's `c=1` q_scaler; `gpumd_init_parameters` re-inits coeffs **and** NN weights uniform(−1,1); L2 (`lambda_2`) shrinks the weights. |
-| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch), is off by default, is per-stage (a stage-1 plateau jumps into Stage 2, surviving resume). |
+| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch), is off by default, is per-stage (a stage-1 plateau jumps into Stage 2, surviving resume). `export_valid_split` reproduces the internal split (verbatim frames, matches `energy_test.out`). |
 | `test_stream_mode.py` | `StreamDataStore` (the training data store): `collate` is bit-exact vs an independently assembled reference (concatenation + offsets + basis straight from the ops functions); metadata/mask consistency incl. missing channels; prefetched `iter_collated` matches direct collate; 2-rank DDP same-seed reproducibility. The DDP case is local-only: set `TORCHNEP_TEST_DDP=1` (skipped in CI — multi-process rendezvous is unreliable on shared runners). |
 | `test_compiled_autograd.py` | `CompiledAutogradForce` (make_fx-materialized autograd forces): outputs and parameter gradients (second-order path through the force loss) match eager autograd across batch shapes on one dynamic graph; energy-only calls fall back to eager. CUDA-only — auto-skipped on CPU hosts/CI. |
 

tests/test_run_seed_and_valid.py

@@ -358,3 +358,51 @@ def test_resume_preserves_valid_split(tmp_path):
     expected = np.array([frames[i]["energy"] / len(frames[i]["species"])
                          for i in val_idx])
     assert np.allclose(ref_resumed, expected, atol=1e-8)
+
+
+def test_export_valid_split_matches_train_nep(tmp_path):
+    """export_valid_split writes GPUMD-ready train/test files that
+    reproduce exactly the valid_ratio split train_nep uses: the reference
+    energies in energy_test.out equal the exported test.xyz energies (in
+    order), the two files partition the input verbatim, and the export is
+    deterministic."""
+    from torchnep import export_valid_split
+    nepin, xyz = _write_run_files(tmp_path)
+
+    out = tmp_path / "run"
+    train_nep(config_file=nepin, data_file=xyz, output_dir=str(out),
+              device="cpu", precision="float64", print_interval=100,
+              restart=False, checkpoint_interval=1000,
+              prediction_interval=1000, run_seed=11, valid_ratio=0.25)
+
+    tp, vp, n_tr, n_va = export_valid_split(
+        xyz, valid_ratio=0.25, run_seed=11,
+        output_dir=str(tmp_path / "split"))
+
+    orig = read_xyz(xyz)
+    tr, va = read_xyz(tp), read_xyz(vp)
+    assert n_tr == len(tr) and n_va == len(va)
+    assert len(tr) + len(va) == len(orig)
+
+    # Reference per-atom energies in energy_test.out (col 2) must equal the
+    # exported test.xyz frames, same order.
+    rows = np.loadtxt(out / "energy_test.out", ndmin=2)
+    assert rows.shape[0] == len(va)
+    ref = rows[:, 1]
+    exp = np.array([f["energy"] / f["natoms"] for f in va])
+    assert np.allclose(ref, exp, atol=1e-10)
+
+    # Verbatim: re-reading the two exports and the original gives the same
+    # frame multiset (match on energy + natoms fingerprints).
+    def fp(frames):
+        return sorted((f["natoms"], round(float(f["energy"]), 10))
+                      for f in frames)
+    assert fp(tr) + fp(va) != []  # sanity
+    assert sorted(fp(tr) + fp(va)) == fp(orig)
+
+    # Deterministic: exporting again is byte-identical.
+    tp2, vp2, _, _ = export_valid_split(
+        xyz, valid_ratio=0.25, run_seed=11,
+        output_dir=str(tmp_path / "split2"))
+    assert open(tp).read() == open(tp2).read()
+    assert open(vp).read() == open(vp2).read()

torchnep/__init__.py

@@ -16,7 +16,9 @@
 from .predict import predict_dataset
 from .train import train_nep
 from .train_sharded import train_nep_sharded
+from .data import export_valid_split
 
 __all__ = [
     "predict_dataset", "train_nep", "train_nep_sharded",
+    "export_valid_split",
 ]

torchnep/data.py

@@ -462,3 +462,61 @@ def build_neighbor_list_np(positions, cell, cutoff):
     return (np.concatenate(all_i).astype(np.int64),
             np.concatenate(all_j).astype(np.int64),
             np.concatenate(all_rij))
+
+
+def valid_split_indices(n_frames: int, valid_ratio: float, run_seed: int):
+    """Train/validation split indices — the exact split ``train_nep`` makes.
+
+    ``train_nep(valid_ratio=r, run_seed=s)`` holds out
+    ``max(1, round(r * n))`` frames drawn from a dedicated torch generator
+    seeded with ``run_seed``. This helper is that draw, factored out so the
+    trainers and :func:`export_valid_split` can never disagree.
+
+    Returns ``(train_idx, valid_idx)`` — both sorted in input-file order.
+    """
+    import torch
+    if run_seed is None:
+        raise ValueError("run_seed is required: the split is drawn from it "
+                         "(train_nep uses the same seed to reproduce it)")
+    if not 0.0 < valid_ratio < 1.0:
+        raise ValueError(f"valid_ratio must be in (0, 1), got {valid_ratio}")
+    g = torch.Generator()
+    g.manual_seed(run_seed)
+    perm = torch.randperm(n_frames, generator=g).tolist()
+    n_val = max(1, int(round(valid_ratio * n_frames)))
+    if n_val >= n_frames:
+        raise ValueError(f"valid_ratio={valid_ratio} leaves no training "
+                         f"frames ({n_frames} total)")
+    val_set = set(perm[:n_val])
+    train_idx = [i for i in range(n_frames) if i not in val_set]
+    return train_idx, sorted(val_set)
+
+
+def export_valid_split(data_file: str, valid_ratio: float, run_seed: int,
+                       output_dir: str = "split"):
+    """Write GPUMD-ready ``train.xyz`` / ``test.xyz`` with train_nep's split.
+
+    Reproduces exactly the validation split that
+    ``train_nep(data_file, valid_ratio=r, run_seed=s)`` uses internally, so
+    the exported pair can train the SAME data partition in GPUMD (or any
+    other code) and loss curves stay comparable. Frames are copied verbatim
+    (raw text, untouched fields and precision), in input-file order.
+
+    Returns ``(train_path, test_path, n_train, n_valid)``.
+    """
+    import os
+    with open(data_file) as f:
+        blocks = _split_frames(f.readlines())
+    train_idx, val_idx = valid_split_indices(len(blocks), valid_ratio,
+                                             run_seed)
+    os.makedirs(output_dir, exist_ok=True)
+    train_path = os.path.join(output_dir, "train.xyz")
+    test_path = os.path.join(output_dir, "test.xyz")
+    src = os.path.abspath(data_file)
+    for path, idxs in ((train_path, train_idx), (test_path, val_idx)):
+        if os.path.abspath(path) == src:
+            raise ValueError(f"output would overwrite the input: {src}")
+        with open(path, "w") as out:
+            for k in idxs:
+                out.writelines(blocks[k])
+    return train_path, test_path, len(train_idx), len(val_idx)

torchnep/train.py

@@ -41,7 +41,7 @@
 from torch.optim.swa_utils import AveragedModel
 
 from .model import NEPModel, slim_model, gpumd_init_parameters
-from .data import read_xyz, parse_nep_in, build_neighbor_list_np
+from .data import read_xyz, parse_nep_in, valid_split_indices, build_neighbor_list_np
 from . import ops
 from . import __version__
 from .predict import predict_from_store
@@ -1324,22 +1324,14 @@ def _log(msg=""):
         _log(f"  read {len(valid_frames)} validation structures "
              f"from {valid_file}")
     elif valid_ratio is not None:
-        if not 0.0 < valid_ratio < 1.0:
-            raise ValueError(f"valid_ratio must be in (0, 1), "
-                             f"got {valid_ratio}")
-        g = torch.Generator()
-        g.manual_seed(run_seed)
-        perm = torch.randperm(len(frames), generator=g).tolist()
-        n_val = max(1, int(round(valid_ratio * len(frames))))
-        if n_val >= len(frames):
-            raise ValueError(f"valid_ratio={valid_ratio} leaves no "
-                             f"training frames ({len(frames)} total)")
-        val_idx = sorted(perm[:n_val])
-        val_set = set(val_idx)
+        # The same draw export_valid_split reproduces (see data.py) — the
+        # split can be exported as GPUMD-ready train.xyz/test.xyz files.
+        train_idx, val_idx = valid_split_indices(len(frames), valid_ratio,
+                                                 run_seed)
         valid_frames = [frames[i] for i in val_idx]
-        frames = [f for i, f in enumerate(frames) if i not in val_set]
-        _log(f"  valid_ratio={valid_ratio}: held out {n_val} frames for "
-             f"validation, {len(frames)} remain for training "
+        frames = [frames[i] for i in train_idx]
+        _log(f"  valid_ratio={valid_ratio}: held out {len(val_idx)} frames "
+             f"for validation, {len(frames)} remain for training "
              f"(split drawn from run_seed)")
 
     # Single-GPU: per-epoch shuffle is done at iteration time via

torchnep/train_sharded.py

@@ -42,7 +42,7 @@
 
 import torch.nn as nn
 from .model import NEPModel, gpumd_init_parameters
-from .data import read_xyz, parse_nep_in
+from .data import read_xyz, parse_nep_in, valid_split_indices
 from . import ops
 from . import __version__
 from .predict import predict_from_store_sharded
@@ -423,22 +423,13 @@ def _log(msg=""):
         _log(f"  read {len(valid_frames)} validation structures "
              f"from {valid_file}")
     elif valid_ratio is not None:
-        if not 0.0 < valid_ratio < 1.0:
-            raise ValueError(f"valid_ratio must be in (0, 1), "
-                             f"got {valid_ratio}")
-        vg = torch.Generator()
-        vg.manual_seed(run_seed)
-        vperm = torch.randperm(len(frames), generator=vg).tolist()
-        n_val = max(1, int(round(valid_ratio * len(frames))))
-        if n_val >= len(frames):
-            raise ValueError(f"valid_ratio={valid_ratio} leaves no "
-                             f"training frames ({len(frames)} total)")
-        val_idx = sorted(vperm[:n_val])
-        val_set = set(val_idx)
+        # Same draw as export_valid_split / train_nep (see data.py).
+        train_idx, val_idx = valid_split_indices(len(frames), valid_ratio,
+                                                 run_seed)
         valid_frames = [frames[i] for i in val_idx]
-        frames = [f for i, f in enumerate(frames) if i not in val_set]
-        _log(f"  valid_ratio={valid_ratio}: held out {n_val} frames for "
-             f"validation, {len(frames)} remain for training "
+        frames = [frames[i] for i in train_idx]
+        _log(f"  valid_ratio={valid_ratio}: held out {len(val_idx)} frames "
+             f"for validation, {len(frames)} remain for training "
              f"(split drawn from run_seed)")
     n_total = len(frames)