comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-29 06:35:28 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

PAO-ML: Stabilize test and add TZV2P basis set

Browse source
This commit is contained in:
Ole Schütt 2025-01-08 22:13:53 +01:00 • committed by Ole Schütt
parent 60e7a04d14
commit 1e46fa001c
2 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
tools/docker/scripts/test_misc.sh
View file

@ -35,7 +35,7 @@ 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-1 --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

11
tools/pao-ml/pao/model.py
View file

@ -51,12 +51,15 @@ class PaoModel(torch.nn.Module):
self.cutoff = cutoff
# Irreps of primary basis
assert prim_basis_name == "DZVP-MOLOPT-GTH" # TODO support more basis sets
# TODO: Export the specs directly from cp2k as part of the .pao files.
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
"DZVP-MOLOPT-GTH/H": "2x0e + 1x1o", # two s-shells, one p-shell
"DZVP-MOLOPT-GTH/O": "2x0e + 2x1o + 1x2e", # two s, two p, one d-shell
"TZV2P-MOLOPT-GGA-GTH-q1/H": "3x0e + 2x1o + 1x2e",
"TZV2P-MOLOPT-GGA-GTH-q6/O": "3x0e + 3x1o + 2x2e + 1x3o",
}
prim_basis_irreps = e3nn.o3.Irreps(prim_basis_specs[kind_name])
basis_specs_key = f"{prim_basis_name}/{kind_name}"
prim_basis_irreps = e3nn.o3.Irreps(prim_basis_specs[basis_specs_key])
assert self.prim_basis_size == prim_basis_irreps.dim
# auxiliary Hamiltonian