PAO-ML: Add scripts for training an equivariant model

tools/docker/scripts/install_misc.sh

@@ -31,6 +31,9 @@ pip3 install --quiet \
   matplotlib \
   requests \
   types-requests \
+  torch \
+  e3nn \
+  scipy \
   mypy==1.5.1
 
 # download inputs for minimax_to_fortran_source.py

tools/docker/scripts/test_misc.sh

@@ -32,6 +32,14 @@ run_test ./tools/precommit/format_fortran_test.py
 run_test ./tools/minimax_tools/minimax_to_fortran_source.py --check
 run_test ./tools/docker/generate_dockerfiles.py --check
 
+# Test pao-ml training.
+run_test ./tools/pao-ml/pao-train.py --kind=H --epochs=200 ./tools/pao-ml/example.pao
+run_test ./tools/pao-ml/pao-retrain.py --model="DZVP-MOLOPT-GTH-PAO4-H.pt" --epochs=200 ./tools/pao-ml/example.pao
+run_test ./tools/pao-ml/pao-validate.py --threshold=1e-2 --model="DZVP-MOLOPT-GTH-PAO4-H.pt" ./tools/pao-ml/example.pao
+run_test ./tools/pao-ml/pao-validate.py --threshold=1e-6 --model="tests/QS/regtest-pao-5/DZVP-MOLOPT-GTH-PAO4-H.pt" ./tools/pao-ml/example.pao
+run_test ./tools/pao-ml/pao-validate.py --threshold=1e-5 --model="tests/QS/regtest-pao-5/DZVP-MOLOPT-GTH-PAO4-O.pt" ./tools/pao-ml/example.pao
+
+run_test mypy --strict ./tools/pao-ml/
 run_test mypy --strict ./tools/minimax_tools/minimax_to_fortran_source.py
 run_test mypy --strict ./tools/dashboard/generate_dashboard.py
 run_test mypy --strict ./tools/dashboard/generate_regtest_survey.py

tools/pao-ml/README.md

@@ -0,0 +1,45 @@
+# PAO-ML training
+
+These are the scripts for training an equivariant PAO-ML model. To create a new model, follow these
+steps:
+
+1. Create and activate a [virtual Python environment](https://docs.python.org/3/tutorial/venv.html):
+
+   ```
+   python3 -m venv ./venv
+   source ./venv/bin/activate
+   ```
+
+1. Install the required Python packages:
+
+   ```
+   pip3 install torch e3nn scipy
+   ```
+
+1. Train a new model:
+
+   ```
+   ./pao-train.py --kind=H training_data1.pao training_data2.pao...
+   ```
+
+1. Retrain the model on some more taining data:
+
+   ```
+   ./pao-retrain.py --model="DZVP-MOLOPT-GTH-PAO4-H.pt" training_data100.pao training_data101.pao...
+   ```
+
+1. Validate the model against test data:
+
+   ```
+   ./pao-validate.py --model="DZVP-MOLOPT-GTH-PAO4-H.pt" test_data1.pao test_data2.pao...
+   ```
+
+1. Use the model in a CP2K run by setting the `PAO_MODEL_FILE` keyword in the kind section:
+
+   ```
+   &KIND H
+     BASIS_SET DZVP-MOLOPT-GTH
+     PAO_BASIS_SIZE 4
+     PAO_MODEL_FILE DZVP-MOLOPT-GTH-PAO4-H.pt
+   &END KIND
+   ```

tools/pao-ml/example.pao

@@ -0,0 +1,30 @@
+Version     4
+Energy            -7.8028963840
+Step     425
+Parametrization     EQUIVARIANT
+Nkinds     2
+Kind              1 O   8
+NParams              1  52
+PrimBasis              1         13 DZVP-MOLOPT-GTH
+PaoBasis              1   4
+NPaoPotentials              1   0
+Kind              2 H   1
+NParams              2  20
+PrimBasis              2          5 DZVP-MOLOPT-GTH
+PaoBasis              2   4
+NPaoPotentials              2   0
+Cell        8.0000000000000000        0.0000000000000000        0.0000000000000000        0.0000000000000000        8.0000000000000000        0.0000000000000000        0.0000000000000000        0.0000000000000000        8.0000000000000000     
+Natoms     6
+Atom               1     O    2.5270000000000006        4.1864999999999997        4.3500000000000005     
+Atom               2     H    2.0350000000000001        4.7115000000000009        3.7100000000000004     
+Atom               3     H    2.1940000000000000        3.2885000000000004        4.2560000000000002     
+Atom               4     O    5.4400000000000004        4.2344999999999997        3.6349999999999998     
+Atom               5     H    5.9649999999999999        4.5744999999999996        4.3650000000000002     
+Atom               6     H    4.5220000000000002        4.2535000000000007        3.9470000000000001     
+Xblock          1  -0.97219572019049405       0.16334145598298072        5.3085713440749652E-002  0.10403752898882876       0.11584659160824060       -7.9301818328872190E-003  -1.5410686264041957E-002  -1.8492276778365915E-002  -2.5342783894250114E-003   4.5102651891888151E-003   4.1411341209456999E-003  -1.0566868220105344E-002   6.2498219102537639E-003   2.3472086470191186E-004  -1.1121087087946725E-005  0.92525105902053273      -0.35449227843300407      -0.10286389296824518        6.6125760590242194E-002  -2.5215356181470038E-002  -7.2477037271285513E-003  -2.8805916640161848E-002  -2.3896643464416154E-002   1.9609923155419231E-002   1.4924492853919646E-002   1.8106821865490712E-002  0.16119977350577630        7.6544516487885988E-002  0.30621068416288677       0.59505620971826734       0.70362928481461118        4.8284458300974498E-002   9.2876243015067611E-002  0.11235870570968362       -1.2916016019097231E-002  -7.5092802140602679E-003  -1.5689336249998843E-003  -3.0136191631121897E-002  -3.8561896986622187E-003  -3.9767642830737440E-003   1.9131247695298144E-003 -0.19043986328178980      -0.69063968602730585       0.66702152075022791       -3.9471751509644146E-002 -0.14308952806313241       0.13796320026960923       -2.6760919619490727E-003   9.5911623766590428E-003   2.0328978650105391E-002  -1.8868433005201973E-004  -1.3418378184515967E-002
+Xblock          2   0.88301186334478610       0.45858285192934001       -6.9292200255783420E-002   5.4665627876413728E-002   3.9323694954076616E-002 -0.12313872312742918       0.47561698021895482       0.59206325043736885      -0.51660010721798288      -0.37467709246612074       -9.7753956676590148E-002  0.16705057309484650       0.62991018854615599       0.30273561616771560       0.68834548506093696       -3.2416256183331112E-002   6.0103059409002998E-002  0.28227603300475479       0.74941310410254691      -0.59498091180555457     
+Xblock          3  -0.88321878793377229      -0.45817005144709744       -9.4014801632015643E-002   7.9733003416048165E-003  -2.1118079208337505E-002  0.12069443763019999      -0.47216689262402184       0.84509711016343125       -3.4077630261589345E-002  0.21522596950825479       -9.8263713501084649E-002  0.16851585867923108      -0.12460817861752686       0.59172975798233352       0.77200549985511580       -3.0812976204735713E-002   5.7344911859215600E-002  0.17543745795742788       0.79038237961264168      -0.58330568092670887     
+Xblock          4  -0.97155934468935978       0.16687517496968285       -5.1077108655797301E-002 -0.14723152083065796        5.2373786610739007E-002   7.9956889487979160E-003   2.3687925690928376E-002  -1.0657708249332962E-002  -3.2898481892559063E-003  -7.6886317241332824E-003  -3.4792165945211713E-003  -1.8309478489571844E-003  -1.0025118343278560E-002   6.9139689455150841E-003   1.1326436440232950E-002 -0.21023593098238513      -0.27742671846974493      -0.93226812661061309       -1.5582185526685893E-002  -1.7805192276814628E-002  -7.9349963948795343E-002  -1.1380467113966777E-002  -1.3874661043664097E-002  -2.0507249201791346E-002  -3.7338327249906933E-002   1.6834051502464621E-002 -0.15872854681481097       -7.0263473018857131E-002  0.29759729726821565       0.86400577871127027      -0.32665602633702051        4.8143888833398887E-002  0.14225023091062589       -6.1330150907657872E-002  -2.5726875774003161E-003   1.9638881001844480E-002   2.2818365435105636E-002  -1.3357177821429174E-002   7.9415262716628025E-003  -3.0240004605119073E-005  -9.2470540726806796E-005 -0.90754117486117702       0.35049397258257026       0.10035936160943225      -0.19164051193710518        7.3968432215091556E-002   2.1192234962159599E-002   8.3101753947402348E-003  -1.9464303177492231E-002   1.5005625299107612E-002  -4.1192281105380442E-004   6.9945130302283545E-003
+Xblock          5   0.88243976900743171       0.46045049804512439       -2.9562869236820971E-002  -5.7229663220037502E-002  -6.6980092692184634E-002  0.11655061088681479      -0.45458128824469879      -0.28939163546582264      -0.57530242703947776      -0.60362017642861199       -9.1310651661635461E-002  0.16261915019816123      -0.17242836689907981      -0.65205658436563829       0.71417736194734749        2.1452320246014792E-003  -4.3910582687340044E-003  0.92762242312444010      -0.35872433987395558      -0.10396778535276149     
+Xblock          6  -0.88730378100476437      -0.45302566667134014       -6.0684066340663664E-003   7.7833688434071529E-003  -8.2319682399070837E-002  -3.5882846691836966E-002 -0.17773219927781272        5.7340877488743468E-002 -0.13206079892228445       0.97251724748914159       0.14429918109724846      -0.26828204648464404       0.34947359057027427       0.88147127762860200        8.7079074712687082E-002  -6.7615486102373993E-004   1.5986990840324965E-003 -0.92805202359323302       0.35832466374049732       0.10158553346294738     
+THE_END

tools/pao-ml/pao-retrain.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+# author: Ole Schuett
+
+import argparse
+import torch
+
+from e3nn import o3  # type: ignore
+from pathlib import Path
+from torch.utils.data import DataLoader
+
+from pao.model import PaoModel
+from pao.dataset import PaoDataset
+from pao.training import train_model
+
+
+# ======================================================================================
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Re-trains an exiting PAO-ML model.")
+    parser.add_argument("--epochs", type=int, default=5000)
+    parser.add_argument("--batch", type=int, default=64)
+    parser.add_argument("--model", type=Path, required=True)
+
+    # Training data files are passed as positional arguments.
+    parser.add_argument("training_data", type=Path, nargs="+")
+    args = parser.parse_args()
+
+    # Load existing model and ignore most cmd arguments.
+    model = torch.jit.load(args.model)  # type: ignore
+    assert model.pao_model_version >= 1
+    print(f"Loaded pre-trained model from file: {args.model}")
+
+    # Load the training data.
+    dataset = PaoDataset(
+        kind_name=model.kind_name,
+        num_neighbors=model.num_neighbors,
+        files=args.training_data,
+    )
+    dataloader = DataLoader(dataset, batch_size=args.batch, shuffle=True)
+    print(f"Found {len(dataset)} training samples of kind '{model.kind_name}'.")
+
+    # Check compatability between model and training data.
+    assert dataset.kind.atomic_number == model.atomic_number
+    assert dataset.kind.pao_basis_size == model.pao_basis_size
+    assert dataset.kind.prim_basis_name == model.prim_basis_name
+    assert dataset.kind.prim_basis_size == model.prim_basis_size
+    assert dataset.feature_kind_names == model.feature_kind_names
+
+    # Train the model.
+    train_model(model, dataloader, args.epochs)
+
+    # Save the model.
+    model.save(args.model)
+    print(f"Saved model to file: {args.model}")
+
+
+# ======================================================================================
+main()
+
+# EOF

tools/pao-ml/pao-train.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+
+# author: Ole Schuett
+
+import torch
+import argparse
+from typing import cast
+from pathlib import Path
+from torch.utils.data import DataLoader
+
+from pao.model import PaoModel
+from pao.dataset import PaoDataset, PaoRecord
+from pao.training import train_model
+
+
+# ======================================================================================
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Trains an equivariant PAO-ML model.")
+    parser.add_argument("--kind", required=True)
+    parser.add_argument("--epochs", type=int, default=5000)
+    parser.add_argument("--batch", type=int, default=64)
+    parser.add_argument("--model", type=Path, default=None)
+
+    # Hyper-parameters - TODO tune default values
+    parser.add_argument("--neighbors", type=int, default=5)
+    parser.add_argument("--distances", type=int, default=10)
+    parser.add_argument("--layers", type=int, default=16)
+    parser.add_argument("--cutoff", type=float, default=6.0)
+
+    # Training data files are passed as positional arguments.
+    parser.add_argument("training_data", type=Path, nargs="+")
+    args = parser.parse_args()
+
+    # Load the training data.
+    dataset = PaoDataset(
+        kind_name=args.kind, num_neighbors=args.neighbors, files=args.training_data
+    )
+    dataloader = DataLoader(dataset, batch_size=args.batch, shuffle=True)
+    print(f"Found {len(dataset)} training samples of kind '{args.kind}'.")
+
+    # Construct the model.
+    model_py = PaoModel(
+        kind_name=dataset.kind.name,
+        atomic_number=dataset.kind.atomic_number,
+        prim_basis_name=dataset.kind.prim_basis_name,
+        prim_basis_size=dataset.kind.prim_basis_size,
+        pao_basis_size=dataset.kind.pao_basis_size,
+        feature_kind_names=dataset.feature_kind_names,
+        num_neighbors=args.neighbors,
+        num_distances=args.distances,
+        num_layers=args.layers,
+        cutoff=args.cutoff,
+    )
+
+    # Compile the model to TorchScript.
+    model = cast(PaoModel, torch.jit.script(model_py))
+
+    num_model_params = sum(p.numel() for p in model.parameters())
+    print(f"PAO-ML model will have {num_model_params} parameters.")
+
+    # Train the model.
+    train_model(model, dataloader, args.epochs)
+
+    # Save the model.
+    default_fn = f"{dataset.kind.prim_basis_name}-PAO{dataset.kind.pao_basis_size}-{args.kind}.pt"
+    fn = args.model or default_fn
+    model.save(fn)
+    print(f"Saved model to file: {fn}")
+
+
+# ======================================================================================
+main()
+
+# EOF

tools/pao-ml/pao-validate.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+
+# author: Ole Schuett
+
+import torch
+import argparse
+import numpy as np
+from e3nn import o3  # type: ignore
+from pathlib import Path
+from torch.utils.data import DataLoader
+
+from pao.model import PaoModel
+from pao.dataset import PaoDataset
+from pao.training import loss_function
+
+
+# ======================================================================================
+def main() -> None:
+    description = "Validates a given equivariant PAO-ML model against test data."
+    parser = argparse.ArgumentParser(description=description)
+    # parser.add_argument("--batch", type=int, default=64)
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--threshold", type=float, default=1.0)
+
+    # Test data files are passed as positional arguments.
+    parser.add_argument("test_data", type=Path, nargs="+")
+    args = parser.parse_args()
+
+    # Load model.
+    model = torch.jit.load(args.model)  # type: ignore
+    assert model.pao_model_version >= 1
+    print(f"Loaded model from file: {args.model}")
+
+    # Load the test data.
+    dataset = PaoDataset(
+        kind_name=model.kind_name,
+        num_neighbors=model.num_neighbors,
+        files=args.test_data,
+    )
+    print(f"Found {len(dataset)} test samples of kind '{model.kind_name}'.")
+
+    # Check compatability between model and test data.
+    assert dataset.kind.atomic_number == model.atomic_number
+    assert dataset.kind.pao_basis_size == model.pao_basis_size
+    assert dataset.kind.prim_basis_name == model.prim_basis_name
+    assert dataset.kind.prim_basis_size == model.prim_basis_size
+    assert dataset.feature_kind_names == model.feature_kind_names
+
+    # Compute losses.
+    losses = []
+    for neighbors_relpos, neighbors_features, label in dataset:
+        inputs = {
+            "neighbors_relpos": neighbors_relpos,
+            "neighbors_features": neighbors_features,
+        }
+        outputs = model(inputs)
+        loss = loss_function(outputs["xblock"], label)
+        losses.append(loss.item())
+
+    print("minimum loss: {:.8e}".format(np.amin(losses).item()))
+    print("median  loss: {:.8e}".format(np.median(losses).item()))
+    print("maximum loss: {:.8e}".format(np.amax(losses).item()))
+
+    assert np.median(losses).item() < args.threshold
+
+
+# ======================================================================================
+main()
+
+# EOF

tools/pao-ml/pao/dataset.py

@@ -0,0 +1,77 @@
+# author: Ole Schuett
+
+import torch
+import numpy as np
+import numpy.typing as npt
+import scipy.spatial  # type: ignore
+from pathlib import Path
+from typing import Any, List, Tuple, Iterator
+from torch.utils.data import Dataset
+
+from .io import parse_pao_file
+
+PaoRecord = Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+
+
+# ======================================================================================
+def _as_tensor(x: List[Any] | npt.NDArray[np.float64]) -> torch.Tensor:
+    return torch.tensor(np.array(x, dtype=np.float32))
+
+
+# ======================================================================================
+class PaoDataset(Dataset[PaoRecord]):
+    def __init__(self, kind_name: str, num_neighbors: int, files: List[Path]):
+        self.neighbors_relpos: List[torch.Tensor] = []
+        self.neighbors_features: List[torch.Tensor] = []
+        self.labels: List[torch.Tensor] = []
+
+        # Load kinds from the first training data file.
+        kinds = parse_pao_file(files[0]).kinds
+        self.kind = kinds[kind_name]
+        self.feature_kind_names = sorted(kinds.keys())
+
+        # Load all training data files.
+        for fn in files:
+            f = parse_pao_file(fn)
+
+            # Build  k-d tree of atom positions.
+            assert 0 < num_neighbors < f.coords.shape[0]
+            assert np.all(f.cell == np.diag(np.diagonal(f.cell)))
+            boxsize = np.diagonal(f.cell)
+            kdtree = scipy.spatial.KDTree(np.mod(f.coords, boxsize), boxsize=boxsize)
+            # alternative: https://wiki.fysik.dtu.dk/ase/ase/neighborlist.html
+
+            for i, k in enumerate(f.atom2kind):
+                if k == kind_name:
+                    # Find indicies of neighbor atoms.
+                    nearest = kdtree.query(f.coords[i], num_neighbors + 1)[1]
+                    neighbors = [j for j in nearest if j != i]  # filter central atom
+
+                    # Compute relative positions of neighbor atoms.
+                    relpos = [f.coords[j] - f.coords[i] for j in neighbors]
+                    self.neighbors_relpos.append(_as_tensor(relpos))
+
+                    # Features of neighbor atoms is the one-hot encoding of their kind.
+                    features = [
+                        f.atom2kind[j] == np.array(self.feature_kind_names)
+                        for j in neighbors
+                    ]
+                    self.neighbors_features.append(_as_tensor(features))
+
+                    # Orthonormalize labels as it's required for the loss_functions.
+                    label = np.linalg.svd(f.xblocks[i], full_matrices=False)[2]
+                    self.labels.append(_as_tensor(label))
+
+    def __len__(self) -> int:
+        return len(self.labels)
+
+    def __getitem__(self, i: int) -> PaoRecord:
+        return self.neighbors_relpos[i], self.neighbors_features[i], self.labels[i]
+
+    # To make mypy happy.
+    def __iter__(self) -> Iterator[PaoRecord]:
+        for i in range(len(self)):
+            yield self[i]
+
+
+# EOF

tools/pao-ml/pao/io.py

@@ -0,0 +1,147 @@
+# author: Ole Schuett
+
+import re
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple
+import numpy as np
+import numpy.typing as npt
+
+KindName = str
+NDArray = npt.NDArray[np.float64]
+
+
+# ======================================================================================
+@dataclass
+class AtomicKind:
+    name: str
+    atomic_number: int
+    nparams: int = -1
+    prim_basis_size: int = -1
+    prim_basis_name: str = ""
+    pao_basis_size: int = -1
+
+
+# ======================================================================================
+@dataclass
+class PaoFile:
+    kinds: Dict[KindName, AtomicKind]
+    atom2kind: List[KindName]
+    cell: NDArray
+    coords: NDArray
+    xblocks: List[NDArray]
+
+
+# ======================================================================================
+def parse_pao_file(path: Path) -> PaoFile:
+    assert str(path).endswith(".pao")
+
+    ikind2name = {}  # maps kind index to kind name
+    atom2kind: List[KindName] = []  # maps atom index to kind name
+    kinds: Dict[KindName, AtomicKind] = {}
+    cell: NDArray = np.zeros((3, 3))
+    coords_list: List[List[str]] = []
+    xblocks: List[NDArray] = []
+
+    for line in path.read_text().strip().split("\n"):
+        parts = line.split()
+        if parts[0] == "Parametrization":
+            assert parts[1] == "EQUIVARIANT"
+
+        elif parts[0] == "Kind":
+            ikind = int(parts[1])
+            kind = AtomicKind(name=parts[2], atomic_number=int(parts[3]))
+            kinds[kind.name] = kind
+            ikind2name[ikind] = kind.name
+
+        elif parts[0] == "NParams":
+            ikind = int(parts[1])
+            kinds[ikind2name[ikind]].nparams = int(parts[2])
+
+        elif parts[0] == "PrimBasis":
+            ikind = int(parts[1])
+            kinds[ikind2name[ikind]].prim_basis_size = int(parts[2])
+            kinds[ikind2name[ikind]].prim_basis_name = parts[3]
+
+        elif parts[0] == "PaoBasis":
+            ikind = int(parts[1])
+            kinds[ikind2name[ikind]].pao_basis_size = int(parts[2])
+
+        elif parts[0] == "Cell":
+            cell = np.array(parts[1:], float).reshape(3, 3)
+
+        elif parts[0] == "Atom":
+            atom2kind.append(parts[2])
+            coords_list.append(parts[3:])
+
+        elif parts[0] == "Xblock":
+            xblocks.append(np.array(parts[2:], float))
+
+    # Convert coordinates to numpy array.
+    coords = np.array(coords_list, float)
+
+    # Reshape xblocks.
+    for iatom, kind_name in enumerate(atom2kind):
+        n = kinds[kind_name].prim_basis_size
+        m = kinds[kind_name].pao_basis_size
+        xblocks[iatom] = xblocks[iatom].reshape(m, n)
+
+    return PaoFile(kinds, atom2kind, cell, coords, xblocks)
+
+
+# ======================================================================================
+def write_pao_file(
+    path: Path,
+    kinds: Dict[KindName, AtomicKind],
+    atom2kind: List[KindName],
+    coords: NDArray,
+    xblocks: List[NDArray],
+) -> None:
+    natoms = coords.shape[0]
+    assert coords.shape[1] == 3
+    assert len(xblocks) == natoms
+
+    output = []
+    output.append("Version 4")
+    output.append("Parametrization EQUIVARIANT")
+    output.append(f"Nkinds {len(kinds)}")
+    for ikind, (kind_name, kind) in enumerate(kinds.items()):
+        i = ikind + 1
+        output.append(f"Kind {i} {kind_name} {kind.atomic_number}")
+        output.append(f"NParams {i} {kind.nparams}")
+        output.append(f"PrimBasis {i} {kind.prim_basis_size} {kind.prim_basis_name}")
+        output.append(f"PaoBasis {i} {kind.pao_basis_size}")
+        output.append(f"NPaoPotentials {i} 0")
+    output.append("Cell 8.0 0.0 0.0   0.0 8.0 0.0   0.0 0.0 8.0")
+    output.append(f"Natoms {natoms}")
+
+    for iatom in range(natoms):
+        c = coords[iatom, :]
+        output.append(f"Atom {iatom + 1} {atom2kind[iatom]} {c[0]} {c[1]} {c[2]}")
+
+    for iatom in range(natoms):
+        kind = kinds[atom2kind[iatom]]
+        assert len(xblocks[iatom].shape) == 2
+        assert xblocks[iatom].shape[0] == kind.pao_basis_size
+        assert xblocks[iatom].shape[1] == kind.prim_basis_size
+        x = xblocks[iatom].flatten()
+        y = " ".join(["%f" % i for i in x])
+        output.append(f"Xblock {iatom + 1} {y}")
+
+    output.append("THE_END")
+    path.write_text("\n".join(output))
+
+
+# ======================================================================================
+def read_cp2k_energy(path: Path) -> float:
+    try:
+        content = path.read_text()
+        m = re.search(r"ENERGY\|(.*)", content)
+        assert m
+        return float(m.group(1).split()[-1])
+    except:
+        print(f"error with: {path}")
+    return float("NaN")
+
+
+# EOF

tools/pao-ml/pao/model.py

@@ -0,0 +1,200 @@
+# author: Ole Schuett
+
+import torch
+import e3nn  # type: ignore
+import warnings
+from typing import Any, Dict, List
+
+
+# ======================================================================================
+class PaoModel(torch.nn.Module):
+    # Ensure these get picked up as attributes by TorchScript.
+    # https://pytorch.org/docs/stable/generated/torch.jit.Attribute.html
+    pao_model_version: int
+    kind_name: str
+    atomic_number: int
+    prim_basis_name: str
+    prim_basis_size: int
+    pao_basis_size: int
+    feature_kind_names: List[str]
+    num_neighbors: int
+    num_distances: int
+    num_layers: int
+    cutoff: float
+
+    def __init__(
+        self,
+        kind_name: str,
+        atomic_number: int,
+        prim_basis_name: str,
+        prim_basis_size: int,
+        pao_basis_size: int,
+        feature_kind_names: List[str],
+        num_neighbors: int,
+        num_distances: int,
+        num_layers: int,
+        cutoff: float,
+    ):
+        super().__init__()
+        self.pao_model_version = 1
+        self.kind_name = kind_name
+        self.atomic_number = atomic_number
+        self.prim_basis_name = prim_basis_name
+        self.prim_basis_size = prim_basis_size
+        self.pao_basis_size = pao_basis_size
+        self.feature_kind_names = feature_kind_names
+
+        # hyper-parameters
+        self.num_neighbors = num_neighbors
+        self.num_distances = num_distances
+        self.num_layers = num_layers
+        self.cutoff = cutoff
+
+        # Irreps of primary basis
+        assert prim_basis_name == "DZVP-MOLOPT-GTH"  # TODO support more basis sets
+        prim_basis_specs = {
+            "O": "2x0e + 2x1o + 1x2e",  # two s-shells, two p-shells, one d-shell
+            "H": "2x0e + 1x1o",  # two s-shells, one p-shell
+        }
+        prim_basis_irreps = e3nn.o3.Irreps(prim_basis_specs[kind_name])
+        assert self.prim_basis_size == prim_basis_irreps.dim
+
+        # auxiliary Hamiltonian
+        self.matrix = SymmetricMatrix(prim_basis_irreps)
+
+        # Irreps of input features, i.e. the descriptor.
+        self.features_irreps = len(self.feature_kind_names) * e3nn.o3.Irrep("0e")
+        self.features_irreps_dim = self.features_irreps.dim
+
+        # Irreps of Spherical Harmonics used for sensing neighbors.
+        self.sensor_irreps = e3nn.o3.Irreps.spherical_harmonics(
+            lmax=self.matrix.input_irreps.lmax
+        )
+
+        # Tensor Product
+        # TODO: fix warning
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            self.tensor_product = e3nn.o3.FullyConnectedTensorProduct(
+                irreps_in1=self.features_irreps,
+                irreps_in2=self.sensor_irreps,
+                irreps_out=self.matrix.input_irreps,
+                shared_weights=False,
+            )
+
+        # Perceptron
+        # Note ReLu does not work well because many of the distance buckets from soft_one_hot_linspace are zero.
+        self.net = e3nn.nn.FullyConnectedNet(
+            hs=[self.num_distances, self.num_layers, self.tensor_product.weight_numel],
+            act=torch.sigmoid,
+        )
+
+        # CP2K uses the yzx convention, while e3nn uses xyz.
+        # https://docs.e3nn.org/en/stable/guide/change_of_basis.html
+        yzx_to_xyz = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
+        self.D_yzx_to_xyz = prim_basis_irreps.D_from_matrix(yzx_to_xyz)
+
+        # Spherical Harmonics
+        self.spherical_harmonics = e3nn.o3.SphericalHarmonics(
+            self.sensor_irreps,
+            normalize=True,
+            normalization="component",
+        )
+
+        # Used for distance embedding.
+        self.distance_buckets = torch.linspace(0.0, self.cutoff, self.num_distances)
+
+    # ----------------------------------------------------------------------------------
+    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        neighbors_relpos = inputs["neighbors_relpos"]
+        neighbors_features = inputs["neighbors_features"]
+
+        assert neighbors_relpos.shape[-2] == self.num_neighbors
+        assert neighbors_relpos.shape[-1] == 3
+        assert neighbors_features.shape[-2] == self.num_neighbors
+        assert neighbors_features.shape[-1] == self.features_irreps_dim
+
+        neighbors_distances = torch.norm(neighbors_relpos, dim=-1)
+
+        # Corresponds to e3nn.math.soft_one_hot_linspace(neighb_dists, basis="gaussian")
+        bucket_width = self.distance_buckets[1] - self.distance_buckets[0]
+        diff = (neighbors_distances[..., None] - self.distance_buckets) / bucket_width
+        distance_embedding = diff.pow(2).neg().exp().div(1.12)
+        weights = self.net(distance_embedding.mul(self.num_distances**0.5))
+        sensors = self.spherical_harmonics(neighbors_relpos)
+        vec_per_neighbor = self.tensor_product(
+            x=neighbors_features.mul(self.num_neighbors**0.5),
+            y=sensors,
+            weight=weights,
+        )
+        h_aux_vec = vec_per_neighbor.sum(dim=-2).div(self.num_neighbors**0.5)
+        h_aux_matrix = self.matrix(h_aux_vec)
+        u_matrix = torch.linalg.eigh(h_aux_matrix)[1]
+        xblock = u_matrix[..., : self.pao_basis_size].transpose(-2, -1)
+        outputs = {"xblock": xblock @ self.D_yzx_to_xyz}
+        return outputs
+
+
+# ======================================================================================
+def flatten_irreps(irreps: e3nn.o3.Irreps) -> List[e3nn.o3.Irrep]:
+    "Helper function to turn an Irreps object into a list of individual Irrep objects."
+    result = []
+    for mul, ir in irreps:
+        result += mul * [ir]
+    return result
+
+
+# ======================================================================================
+def dim(l: int) -> int:
+    return 2 * l + 1
+
+
+# ======================================================================================
+class SymmetricMatrix(torch.nn.Module):
+    def __init__(self, basis_irreps: Any):
+        super().__init__()
+        self.basis_irreps = basis_irreps
+        self.basis_irreps_ls: List[int] = basis_irreps.ls
+
+        # Compute irreps required to represent a matrix
+        self.input_irreps = e3nn.o3.Irreps()
+        self.wigner_blocks = []
+        for i, a in enumerate(flatten_irreps(basis_irreps)):
+            for j, b in enumerate(flatten_irreps(basis_irreps)):
+                if j > i:
+                    continue  # skip upper triangle
+                self.input_irreps += a * b
+
+                # Pre-compute wigner blocks
+                for lk in range(abs(a.l - b.l), a.l + b.l + 1):
+                    self.wigner_blocks.append(e3nn.o3.wigner_3j(a.l, b.l, lk))
+
+        self.input_irreps_ls = self.input_irreps.ls
+
+    # ----------------------------------------------------------------------------------
+    def forward(self, vector: torch.Tensor) -> torch.Tensor:
+        assert vector.shape[-1] == sum(dim(l) for l in self.input_irreps_ls)
+        basis_size = sum(dim(l) for l in self.basis_irreps_ls)
+        matrix = torch.zeros(vector.shape[:-1] + (basis_size, basis_size))
+        matrix[..., :, :] = torch.eye(basis_size)  # ensure matrix is diagonalizable
+        c = 0  # position in vector
+        z = 0  # position in self.wigner_blocks
+        for i, li in enumerate(self.basis_irreps_ls):
+            a = sum(dim(l) for l in self.basis_irreps_ls[:i])  # first matrix row
+            for j, lj in enumerate(self.basis_irreps_ls):
+                if j > i:
+                    continue  # skip upper triangle
+                b = sum(dim(l) for l in self.basis_irreps_ls[:j])  # first matrix col
+                for lk in range(abs(li - lj), li + lj + 1):
+                    # TODO the wigner blocks are mostly zeros - not sure pytorch takes advantage of that.
+                    wigner_block = self.wigner_blocks[z]
+                    coeffs = vector[..., c : c + dim(lk)]
+                    block = torch.tensordot(coeffs, wigner_block, dims=[[-1], [-1]])
+                    matrix[..., a : a + dim(li), b : b + dim(lj)] += block
+                    c += dim(lk)
+                    z += 1
+
+        return matrix
+
+
+# EOF

tools/pao-ml/pao/training.py

@@ -0,0 +1,41 @@
+# author: Ole Schuett
+
+import torch
+from torch.utils.data import DataLoader
+
+from .model import PaoModel
+from .dataset import PaoDataset, PaoRecord
+
+
+# ======================================================================================
+def loss_function(prediction: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
+    # This assumes the columns of prediction and label are orthonormal.
+    p1 = prediction.transpose(-2, -1) @ prediction
+    p2 = label.transpose(-2, -1) @ label
+    return (p1 - p2).pow(2).mean()
+
+
+# ======================================================================================
+def train_model(
+    model: PaoModel, dataloader: DataLoader[PaoRecord], epochs: int
+) -> None:
+    # Train the model.
+    optim = torch.optim.Adam(model.parameters())
+    for epoch in range(epochs + 1):
+        optim.zero_grad()
+        for neighbors_relpos, neighbors_features, label in dataloader:
+            inputs = {
+                "neighbors_relpos": neighbors_relpos,
+                "neighbors_features": neighbors_features,
+            }
+            outputs = model(inputs)
+            loss = loss_function(outputs["xblock"], label)
+            loss.backward()  # type: ignore
+            optim.step()
+        if epoch % 100 == 0:
+            print(f"epoch: {epoch:5d} | loss: {loss:.8e}")
+
+    print(f"Training complete, final loss: {loss:.8e}")
+
+
+# EOF