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

PAO-ML: Fix NaNs in forces (#5052)

Browse source
This commit is contained in:
Ole Schütt 2026-04-09 19:44:26 +02:00 • committed by GitHub
parent 7d4b1410b6
commit 4c4babd79f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/pao_model.F
View file

@ -361,6 +361,7 @@ CONTAINS
CPASSERT(SIZE(predicted_xblock, 1) == n)
CPASSERT(SIZE(predicted_xblock, 2) == m)
CPASSERT(SIZE(predicted_xblock, 3) == 1)
CPASSERT(ALL(predicted_xblock == predicted_xblock)) ! checking for NaNs
IF (PRESENT(block_X)) THEN
block_X = RESHAPE(predicted_xblock, [n*m, 1])
END IF
@ -375,6 +376,7 @@ CONTAINS
NULLIFY (edge_vectors_grad)
CALL torch_tensor_data_ptr(edge_vectors_grad_tensor, edge_vectors_grad)
CPASSERT(SIZE(edge_vectors_grad, 1) == 3 .AND. SIZE(edge_vectors_grad, 2) == num_edges)
CPASSERT(ALL(edge_vectors_grad == edge_vectors_grad)) ! checking for NaNs
DO iedge = 1, num_edges
jneighbor = INT(edge_index(iedge, 1) + 1)
kneighbor = INT(edge_index(iedge, 2) + 1)

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

@ -75,6 +75,8 @@ class PaoModel(SequentialGraphNetwork): # type: ignore
prim_basis_specs = {
"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
"DZVP-MOLOPT-GGA-GTH-q1/H": "2x0e + 1x1o",
"DZVP-MOLOPT-GGA-GTH-q6/O": "2x0e + 2x1o + 1x2e",
"TZV2P-MOLOPT-GGA-GTH-q1/H": "3x0e + 2x1o + 1x2e",
"TZV2P-MOLOPT-GGA-GTH-q6/O": "3x0e + 3x1o + 2x2e + 1x3o",
}
@ -227,7 +229,8 @@ class SymmetricMatrix(torch.nn.Module):
assert vector.shape[-1] == sum(dim(l) for l in self.irreps_in_ls)
basis_size = sum(dim(l) for l in self.irreps_prim_basis_ls)
matrix = torch.zeros(vector.shape[:-1] + (basis_size, basis_size))
matrix[..., :, :] = torch.eye(basis_size) # ensure matrix is diagonalizable
# Ensure matrix has distinct eigenvalues as needed for grad of torch.linalg.eigh().
matrix[..., :, :] = torch.diag(torch.arange(1, basis_size + 1))
c = 0 # position in vector
z = 0 # position in self.wigner_blocks
for i, li in enumerate(self.irreps_prim_basis_ls):