From b2ac3daf86c7f46463e5ad3d1939fa68c947da1f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ole=20Sch=C3=BCtt?= <ole.schuett@cp2k.org>
Date: Fri, 10 Jul 2026 16:55:19 +0200
Subject: [PATCH] PAO-ML: Add pao-diff.py tool (#5577)

---
 tools/pao-ml/pao-diff.py    | 40 +++++++++++++++++++++++++++++++++++++
 tools/pao-ml/pao/dataset.py | 14 ++++++++-----
 2 files changed, 49 insertions(+), 5 deletions(-)
 create mode 100755 tools/pao-ml/pao-diff.py

diff --git a/tools/pao-ml/pao-diff.py b/tools/pao-ml/pao-diff.py
new file mode 100755
index 0000000000..88a8af844e
--- /dev/null
+++ b/tools/pao-ml/pao-diff.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+# author: Ole Schuett
+
+import numpy as np
+import argparse
+from pathlib import Path
+
+from pao.io import parse_pao_file
+from pao.dataset import prepare_xblock
+from pao.training import loss_function
+
+
+# ======================================================================================
+def main() -> None:
+    description = "Compares two PAO files by computing the loss between their xblocks."
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument("pao_file_a", type=Path)
+    parser.add_argument("pao_file_b", type=Path)
+    args = parser.parse_args()
+
+    a = parse_pao_file(args.pao_file_a)
+    b = parse_pao_file(args.pao_file_b)
+
+    assert a.kinds == b.kinds
+    assert a.atom2kind == b.atom2kind
+    assert np.all(a.cell == b.cell)
+    assert np.all(a.coords == b.coords)
+
+    for iatom in range(len(a.atom2kind)):
+        xa = prepare_xblock(a.xblocks[iatom])
+        xb = prepare_xblock(b.xblocks[iatom])
+        loss = loss_function(xa, xb)
+        print(f"iatom: {iatom+1}  kind: {a.atom2kind[iatom]}  loss: {loss:.8e}")
+
+
+# ======================================================================================
+main()
+
+# EOF
diff --git a/tools/pao-ml/pao/dataset.py b/tools/pao-ml/pao/dataset.py
index 51337a9f58..40d5ed00a8 100644
--- a/tools/pao-ml/pao/dataset.py
+++ b/tools/pao-ml/pao/dataset.py
@@ -29,6 +29,14 @@ def _as_int_tensor(x: List[Any] | npt.NDArray[np.int64]) -> torch.Tensor:
     return torch.tensor(np.array(x, dtype=np.int64))
 
 
+# ======================================================================================
+def prepare_xblock(xblock: npt.NDArray[np.float64]) -> torch.Tensor:
+    # Orthonormalize labels as it's required for the loss_functions.
+    ortho_xblock = np.linalg.svd(xblock, full_matrices=False)[2]
+    # Add an extra leading dimension to simplify the batching.
+    return _as_float_tensor(ortho_xblock).unsqueeze(0)
+
+
 # ======================================================================================
 class PaoDataset(Dataset[AtomicDataDict]):
     def __init__(
@@ -86,17 +94,13 @@ class PaoDataset(Dataset[AtomicDataDict]):
 
                 # TODO remove edges that are more than num_layers hops await from iatom.
 
-                # Orthonormalize labels as it's required for the loss_functions.
-                xblock = np.linalg.svd(f.xblocks[iatom], full_matrices=False)[2]
-
                 # Collect all tensors into an AtomicDataDict for NequIP.
                 self.examples.append(
                     {
                         "atom_types": _as_int_tensor(neighbor_atom_types),
                         "edge_index": _as_int_tensor(edge_index).T,
                         "edge_vectors": _as_float_tensor(edge_vectors),
-                        # Adding an extra leading dimension to simplify the batching.
-                        "xblock": _as_float_tensor(xblock).unsqueeze(0),
+                        "xblock": prepare_xblock(f.xblocks[iatom]),
                         # The "pos" key is used by batched_from_list() to compute edge_index offsets.
                         "pos": torch.zeros(len(neighbor_pos)),
                         # Within an example the central atom is always the first atom.