From 9f0ae1801a50e08f92dcc8688735d2d360a0650f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ole=20Sch=C3=BCtt?= Date: Tue, 7 Jan 2025 21:54:18 +0100 Subject: [PATCH] PAO-ML: Add scripts for training an equivariant model --- tools/docker/scripts/install_misc.sh | 3 + tools/docker/scripts/test_misc.sh | 8 ++ tools/pao-ml/README.md | 45 ++++++ tools/pao-ml/example.pao | 30 ++++ tools/pao-ml/pao-retrain.py | 60 ++++++++ tools/pao-ml/pao-train.py | 74 ++++++++++ tools/pao-ml/pao-validate.py | 70 ++++++++++ tools/pao-ml/pao/__init__.py | 0 tools/pao-ml/pao/dataset.py | 77 +++++++++++ tools/pao-ml/pao/io.py | 147 ++++++++++++++++++++ tools/pao-ml/pao/model.py | 200 +++++++++++++++++++++++++++ tools/pao-ml/pao/training.py | 41 ++++++ 12 files changed, 755 insertions(+) create mode 100644 tools/pao-ml/README.md create mode 100644 tools/pao-ml/example.pao create mode 100755 tools/pao-ml/pao-retrain.py create mode 100755 tools/pao-ml/pao-train.py create mode 100755 tools/pao-ml/pao-validate.py create mode 100644 tools/pao-ml/pao/__init__.py create mode 100644 tools/pao-ml/pao/dataset.py create mode 100644 tools/pao-ml/pao/io.py create mode 100644 tools/pao-ml/pao/model.py create mode 100644 tools/pao-ml/pao/training.py diff --git a/tools/docker/scripts/install_misc.sh b/tools/docker/scripts/install_misc.sh index fa7ec3a694..2450e61bdb 100755 --- a/tools/docker/scripts/install_misc.sh +++ b/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 diff --git a/tools/docker/scripts/test_misc.sh b/tools/docker/scripts/test_misc.sh index e10c870778..8c2bc46ca5 100755 --- a/tools/docker/scripts/test_misc.sh +++ b/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 diff --git a/tools/pao-ml/README.md b/tools/pao-ml/README.md new file mode 100644 index 0000000000..3f610e1bbf --- /dev/null +++ b/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 + ``` diff --git a/tools/pao-ml/example.pao b/tools/pao-ml/example.pao new file mode 100644 index 0000000000..0a8e952dc5 --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao-retrain.py b/tools/pao-ml/pao-retrain.py new file mode 100755 index 0000000000..5b985c12ce --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao-train.py b/tools/pao-ml/pao-train.py new file mode 100755 index 0000000000..14a9397090 --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao-validate.py b/tools/pao-ml/pao-validate.py new file mode 100755 index 0000000000..cb23b55768 --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao/__init__.py b/tools/pao-ml/pao/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/pao-ml/pao/dataset.py b/tools/pao-ml/pao/dataset.py new file mode 100644 index 0000000000..5482dbb8ee --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao/io.py b/tools/pao-ml/pao/io.py new file mode 100644 index 0000000000..199fe2cfb7 --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao/model.py b/tools/pao-ml/pao/model.py new file mode 100644 index 0000000000..d3e524484f --- /dev/null +++ b/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 diff --git a/tools/pao-ml/pao/training.py b/tools/pao-ml/pao/training.py new file mode 100644 index 0000000000..1d9fff430b --- /dev/null +++ b/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