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

Fix indexing and enable stress prediction in NequIP & Allegro (#3428)

Browse source 
Co-authored-by: Maria Bilichenko <mbilichenko@tcopt12.chem.uzh.ch>
This commit is contained in:
Maria Bilichenko 2024-05-24 07:30:00 +02:00 • committed by GitHub
parent 67d2f81faf
commit 519c5118ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 342 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
data/Allegro/gra-water-deployed-neq060sp.pth Normal file
View file

Binary file not shown.

BIN
data/Allegro/si-deployed-neq060dp.pth Normal file
View file

Binary file not shown.

BIN
data/Allegro/si-deployed.pth
View file

Binary file not shown.

BIN
data/Allegro/water-gra-film-double.pth
View file

Binary file not shown.

BIN
data/NequIP/water-deployed-neq060dp.pth Normal file
View file

Binary file not shown.

BIN
data/NequIP/water-deployed-neq060sp.pth Normal file
View file

Binary file not shown.

BIN
data/NequIP/water-double.pth
View file

Binary file not shown.

BIN
data/NequIP/water.pth
View file

Binary file not shown.

2
src/fist_force.F
View file

@ -290,7 +290,7 @@ CONTAINS
! Compute embedding function and manybody energy
CALL energy_manybody(fist_nonbond_env, atomic_kind_set, local_particles, particle_set, &
cell, pot_manybody, para_env, mm_section)
cell, pot_manybody, para_env, mm_section, use_virial)
! Nonbond contribution + Manybody Forces
IF (shell_present) THEN

180
src/force_fields_input.F
View file

@ -799,6 +799,7 @@ CONTAINS
nonbonded%pot(start + n_items)%pot%set(1)%nequip%unit_energy = unit_energy
nonbonded%pot(start + n_items)%pot%set(1)%nequip%unit_cell = unit_cell
CALL read_nequip_data(nonbonded%pot(start + n_items)%pot%set(1)%nequip)
CALL check_cp2k_atom_names_in_torch(atm_names, nonbonded%pot(start + n_items)%pot%set(1)%nequip%type_names_torch)
nonbonded%pot(start + n_items)%pot%rcutsq = nonbonded%pot(start + n_items)%pot%set(1)%nequip%rcutsq
n_items = n_items + 1
END DO
@ -847,11 +848,11 @@ CONTAINS
nonbonded%pot(start + n_items)%pot%set(1)%allegro%unit_energy = unit_energy
nonbonded%pot(start + n_items)%pot%set(1)%allegro%unit_cell = unit_cell
CALL read_allegro_data(nonbonded%pot(start + n_items)%pot%set(1)%allegro)
CALL check_cp2k_atom_names_in_torch(atm_names, nonbonded%pot(start + n_items)%pot%set(1)%allegro%type_names_torch)
nonbonded%pot(start + n_items)%pot%rcutsq = nonbonded%pot(start + n_items)%pot%set(1)%allegro%rcutsq
n_items = n_items + 1
END DO
END DO
END SUBROUTINE read_allegro_section
! **************************************************************************************************
@ -2433,8 +2434,10 @@ CONTAINS
TYPE(nequip_pot_type), POINTER :: nequip
CHARACTER(len=*), PARAMETER :: routineN = 'read_nequip_data'
CHARACTER(LEN=1), PARAMETER :: delimiter = ' '
CHARACTER(LEN=default_path_length) :: allow_tf32_str, config_str, cutoff_str
CHARACTER(LEN=100), ALLOCATABLE, DIMENSION(:) :: tokenized_string
CHARACTER(LEN=default_path_length) :: allow_tf32_str, cutoff_str, types_str
INTEGER :: handle
LOGICAL :: allow_tf32, found
@ -2449,6 +2452,16 @@ CONTAINS
nequip%nequip_version = torch_model_read_metadata(nequip%nequip_file_name, "nequip_version")
cutoff_str = torch_model_read_metadata(nequip%nequip_file_name, "r_max")
types_str = torch_model_read_metadata(nequip%nequip_file_name, "type_names")
CALL tokenize_string(TRIM(types_str), delimiter, tokenized_string)
IF (ALLOCATED(nequip%type_names_torch)) THEN
DEALLOCATE (nequip%type_names_torch)
END IF
ALLOCATE (nequip%type_names_torch(SIZE(tokenized_string)))
nequip%type_names_torch(:) = tokenized_string(:)
READ (cutoff_str, *) nequip%rcutsq
nequip%rcutsq = cp_unit_to_cp2k(nequip%rcutsq, nequip%unit_coords)
nequip%rcutsq = nequip%rcutsq*nequip%rcutsq
@ -2457,9 +2470,18 @@ CONTAINS
nequip%unit_energy_val = cp_unit_to_cp2k(nequip%unit_energy_val, nequip%unit_energy)
nequip%unit_cell_val = cp_unit_to_cp2k(nequip%unit_cell_val, nequip%unit_cell)
! look in config which contains all the .yaml file options to see if we use float32 or float64
config_str = torch_model_read_metadata(nequip%nequip_file_name, "config")
CALL read_default_dtype(config_str, nequip%do_nequip_sp)
IF (torch_model_read_metadata(nequip%nequip_file_name, "default_dtype") == "float32" .AND. &
torch_model_read_metadata(nequip%nequip_file_name, "model_dtype") == "float32") THEN
nequip%do_nequip_sp = .TRUE.
ELSE IF (torch_model_read_metadata(nequip%nequip_file_name, "default_dtype") == "float64" .AND. &
torch_model_read_metadata(nequip%nequip_file_name, "model_dtype") == "float64") THEN
nequip%do_nequip_sp = .FALSE.
ELSE
CALL cp_abort(__LOCATION__, &
"Both default_dtype and model_dtype should be either float32 or float64. Currently, default_dtype is <"// &
torch_model_read_metadata(nequip%nequip_file_name, "default_dtype")//"> and model_dtype is <"// &
torch_model_read_metadata(nequip%nequip_file_name, "model_dtype")//">.")
END IF
allow_tf32_str = torch_model_read_metadata(nequip%nequip_file_name, "allow_tf32")
allow_tf32 = (TRIM(allow_tf32_str) == "1")
@ -2482,8 +2504,10 @@ CONTAINS
TYPE(allegro_pot_type), POINTER :: allegro
CHARACTER(len=*), PARAMETER :: routineN = 'read_allegro_data'
CHARACTER(LEN=1), PARAMETER :: delimiter = ' '
CHARACTER(LEN=default_path_length) :: allow_tf32_str, config_str, cutoff_str
CHARACTER(LEN=100), ALLOCATABLE, DIMENSION(:) :: tokenized_string
CHARACTER(LEN=default_path_length) :: allow_tf32_str, cutoff_str, types_str
INTEGER :: handle
LOGICAL :: allow_tf32, found
@ -2503,6 +2527,15 @@ CONTAINS
"> has not been deployed; did you forget to run `nequip-deploy`?")
END IF
cutoff_str = torch_model_read_metadata(allegro%allegro_file_name, "r_max")
types_str = torch_model_read_metadata(allegro%allegro_file_name, "type_names")
CALL tokenize_string(TRIM(types_str), delimiter, tokenized_string)
IF (ALLOCATED(allegro%type_names_torch)) THEN
DEALLOCATE (allegro%type_names_torch)
END IF
ALLOCATE (allegro%type_names_torch(SIZE(tokenized_string)))
allegro%type_names_torch(:) = tokenized_string(:)
READ (cutoff_str, *) allegro%rcutsq
allegro%rcutsq = cp_unit_to_cp2k(allegro%rcutsq, allegro%unit_coords)
allegro%rcutsq = allegro%rcutsq*allegro%rcutsq
@ -2511,9 +2544,18 @@ CONTAINS
allegro%unit_energy_val = cp_unit_to_cp2k(allegro%unit_energy_val, allegro%unit_energy)
allegro%unit_cell_val = cp_unit_to_cp2k(allegro%unit_cell_val, allegro%unit_cell)
! look in config which contains all the .yaml file options to see if we use float32 or float64
config_str = torch_model_read_metadata(allegro%allegro_file_name, "config")
CALL read_default_dtype(config_str, allegro%do_allegro_sp)
IF (torch_model_read_metadata(allegro%allegro_file_name, "default_dtype") == "float32" .AND. &
torch_model_read_metadata(allegro%allegro_file_name, "model_dtype") == "float32") THEN
allegro%do_allegro_sp = .TRUE.
ELSE IF (torch_model_read_metadata(allegro%allegro_file_name, "default_dtype") == "float64" .AND. &
torch_model_read_metadata(allegro%allegro_file_name, "model_dtype") == "float64") THEN
allegro%do_allegro_sp = .FALSE.
ELSE
CALL cp_abort(__LOCATION__, &
"Both default_dtype and model_dtype should be either float32 or float64. Currently, default_dtype is <"// &
torch_model_read_metadata(allegro%allegro_file_name, "default_dtype")//"> and model_dtype is <"// &
torch_model_read_metadata(allegro%allegro_file_name, "model_dtype")//">.")
END IF
allow_tf32_str = torch_model_read_metadata(allegro%allegro_file_name, "allow_tf32")
allow_tf32 = (TRIM(allow_tf32_str) == "1")
@ -2528,47 +2570,102 @@ CONTAINS
END SUBROUTINE read_allegro_data
! **************************************************************************************************
!> \brief reads the default_dtype used in the Allegro/NequIP model by parsing the config file
!> \param config_str ...
!> \param do_model_sp ...
!> \author Gabriele Tocci
!> \brief returns tokenized string of kinds from .pth file
!> \param element ...
!> \param delimiter ...
!> \param tokenized_array ...
!> \author Maria Bilichenko
! **************************************************************************************************
SUBROUTINE read_default_dtype(config_str, do_model_sp)
SUBROUTINE tokenize_string(element, delimiter, tokenized_array)
CHARACTER(LEN=*), INTENT(IN) :: element
CHARACTER(LEN=1), INTENT(IN) :: delimiter
CHARACTER(LEN=100), ALLOCATABLE, DIMENSION(:), &
INTENT(OUT) :: tokenized_array
CHARACTER(LEN=default_path_length) :: config_str
LOGICAL :: do_model_sp
CHARACTER(LEN=100) :: temp_kinds
INTEGER :: end_pos, i, num_elements, start
LOGICAL, ALLOCATABLE, DIMENSION(:) :: delim_positions
CHARACTER(len=*), PARAMETER :: routineN = 'read_default_dtype'
! Find positions of delimiter within element
ALLOCATE (delim_positions(LEN(element)))
delim_positions = .FALSE.
INTEGER :: handle, i, idx, len_config
DO i = 1, LEN(element)
IF (element(i:i) == delimiter) delim_positions(i) = .TRUE.
END DO
CALL timeset(routineN, handle)
num_elements = COUNT(delim_positions) + 1
len_config = LEN_TRIM(config_str)
idx = INDEX(config_str, "default_dtype:")
IF (idx /= 0) THEN
i = idx + 14 ! skip over "default_dtype:"
DO WHILE (i <= len_config .AND. config_str(i:i) == " ")
i = i + 1 ! skip over any whitespace
END DO
ALLOCATE (tokenized_array(num_elements))
IF (i > len_config) THEN
CALL cp_abort(__LOCATION__, &
"No default_dtype found, check the Nequip/Allegro .yaml or .pth files."// &
" Default_dtype should be either <float32> or <float64>.")
ELSE IF (config_str(i:i + 6) == "float32") THEN
do_model_sp = .TRUE.
ELSE IF (config_str(i:i + 6) == "float64") THEN
do_model_sp = .FALSE.
ELSE
CALL cp_abort(__LOCATION__, &
"The default_dtype should be either <float32> or <float64>."// &
" Check the NequIP/Allegro .yaml and .pth files.")
start = 1
DO i = 1, num_elements
IF (LEN(element) < 3 .AND. COUNT(delim_positions) == 0) THEN ! if there is 1 kind only and it has one or two
!characters (C or Cl) the end_pos will be the index of the last character (1 or 2)
end_pos = LEN(element)
ELSE ! else, the end_pos is determined by the index of the space - 1
end_pos = find_end_pos(start, delim_positions)
END IF
END IF
temp_kinds = element(start:end_pos)
IF (TRIM(temp_kinds) /= '') THEN
tokenized_array(i) = temp_kinds
END IF
start = end_pos + 2
END DO
DEALLOCATE (delim_positions)
END SUBROUTINE tokenize_string
CALL timestop(handle)
END SUBROUTINE read_default_dtype
! **************************************************************************************************
!> \brief finds the position of the atom by the spacing
!> \param start ...
!> \param delim_positions ...
!> \return ...
!> \author Maria Bilichenko
! **************************************************************************************************
INTEGER FUNCTION find_end_pos(start, delim_positions)
INTEGER, INTENT(IN) :: start
LOGICAL, DIMENSION(:), INTENT(IN) :: delim_positions
INTEGER :: end_pos, i
end_pos = start
DO i = start, SIZE(delim_positions)
IF (delim_positions(i)) THEN
end_pos = i - 1
EXIT
END IF
END DO
find_end_pos = end_pos
END FUNCTION find_end_pos
! **************************************************************************************************
!> \brief checks if all the ATOMS from *.inp file are available in *.pth file
!> \param cp2k_inp_atom_types ...
!> \param torch_atom_types ...
!> \author Maria Bilichenko
! **************************************************************************************************
SUBROUTINE check_cp2k_atom_names_in_torch(cp2k_inp_atom_types, torch_atom_types)
CHARACTER(LEN=*), DIMENSION(:), INTENT(IN) :: cp2k_inp_atom_types, torch_atom_types
INTEGER :: i, j
LOGICAL :: found
DO i = 1, SIZE(cp2k_inp_atom_types)
found = .FALSE.
DO j = 1, SIZE(torch_atom_types)
IF (TRIM(cp2k_inp_atom_types(i)) == TRIM(torch_atom_types(j))) THEN
found = .TRUE.
EXIT
END IF
END DO
IF (.NOT. found) THEN
CALL cp_abort(__LOCATION__, &
"Atom "//TRIM(cp2k_inp_atom_types(i))// &
" is defined in the CP2K input file but is missing in the torch model file")
END IF
END DO
END SUBROUTINE check_cp2k_atom_names_in_torch
! **************************************************************************************************
!> \brief reads TABPOT potential from file
@ -2626,5 +2723,4 @@ CONTAINS
CALL cp_print_key_finished_output(iw, logger, mm_section, "PRINT%FF_INFO")
CALL timestop(handle)
END SUBROUTINE read_tabpot_data
END MODULE force_fields_input

48
src/input_cp2k_mm.F
View file

@ -1485,7 +1485,9 @@ CONTAINS
CPASSERT(.NOT. ASSOCIATED(section))
CALL section_create(section, __LOCATION__, name="NEQUIP", &
description="This section specifies the input parameters for NEQUIP potential type "// &
"based on equivariant neural networks with message passing. "// &
"based on equivariant neural networks with message passing. Starting from the NequIP 0.6.0, "// &
"one can predict stress if the config.yaml file has the StressForceOutput keyword, "// &
"regardless of whether the model has been trained on the stress. "// &
"Requires linking with libtorch library from <https://pytorch.org/cppdocs/installing.html>.", &
citations=(/Batzner2022/), n_keywords=1, n_subsections=0, repeats=.FALSE.)
@ -1509,25 +1511,33 @@ CONTAINS
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_COORDS", &
description="Units of coordinates in the NEQUIP model.pth file.", &
description="Units of coordinates in the NEQUIP model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT angstrom", default_c_val="angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_ENERGY", &
description="Units of energy in the NEQUIP model.pth file.", &
usage="UNIT hartree", default_c_val="hartree")
description="Units of energy in the NEQUIP model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT hartree", default_c_val="eV")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_FORCES", &
description="Units of the forces in the NEQUIP model.pth file.", &
usage="UNIT hartree/bohr", default_c_val="hartree/bohr")
description="Units of the forces in the NEQUIP model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT hartree/bohr", default_c_val="eV/Angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_CELL", &
description="Units of the cell vectors in the NEQUIP model.pth file.", &
description="Units of the cell vectors in the NEQUIP model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT angstrom", default_c_val="angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
@ -1547,7 +1557,9 @@ CONTAINS
CPASSERT(.NOT. ASSOCIATED(section))
CALL section_create(section, __LOCATION__, name="ALLEGRO", &
description="This section specifies the input parameters for ALLEGRO potential type "// &
"based on equivariant neural network potentials. "// &
"based on equivariant neural network potentials. Starting from the NequIP 0.6.0, "// &
"one can predict stress if the config.yaml file has the StressForceOutput keyword, "// &
"regardless of whether the model has been trained on the stress. "// &
"Requires linking with libtorch library from <https://pytorch.org/cppdocs/installing.html>.", &
citations=(/Musaelian2023/), n_keywords=1, n_subsections=0, repeats=.FALSE.)
@ -1571,25 +1583,33 @@ CONTAINS
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_COORDS", &
description="Units of coordinates in the ALLEGRO model.pth file.", &
description="Units of coordinates in the ALLEGRO model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT angstrom", default_c_val="angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_ENERGY", &
description="Units of energy in the ALLEGRO model.pth file.", &
usage="UNIT hartree", default_c_val="hartree")
description="Units of energy in the ALLEGRO model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT hartree", default_c_val="eV")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_FORCES", &
description="Units of the forces in the ALLEGRO model.pth file.", &
usage="UNIT hartree/bohr", default_c_val="hartree/bohr")
description="Units of the forces in the ALLEGRO model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT hartree/bohr", default_c_val="eV/Angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)
CALL keyword_create(keyword, __LOCATION__, name="UNIT_CELL", &
description="Units of the cell vectors in the ALLEGRO model.pth file.", &
description="Units of the cell vectors in the ALLEGRO model.pth file. "// &
"The units of positions, energies and forces must be self-consistent: "// &
"e.g. coordinates in Angstrom, energies in eV, forces in eV/Angstrom. ", &
usage="UNIT angstrom", default_c_val="angstrom")
CALL section_add_keyword(section, keyword)
CALL keyword_release(keyword)

68
src/manybody_allegro.F
View file

@ -214,13 +214,15 @@ CONTAINS
!> \param pot_allegro ...
!> \param fist_nonbond_env ...
!> \param unique_list_a ...
!> \param use_virial ...
!> \par History
!> Implementation of the allegro potential - [gtocci] 2023
!> Index mapping of atoms from .xyz to Allegro config.yaml file - [mbilichenko] 2024
!> \author Gabriele Tocci - University of Zurich
! **************************************************************************************************
SUBROUTINE allegro_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, &
potparm, allegro, glob_loc_list_a, r_last_update_pbc, &
pot_allegro, fist_nonbond_env, unique_list_a)
pot_allegro, fist_nonbond_env, unique_list_a, use_virial)
TYPE(fist_neighbor_type), POINTER :: nonbonded
TYPE(particle_type), POINTER :: particle_set(:)
@ -233,12 +235,13 @@ CONTAINS
REAL(kind=dp) :: pot_allegro
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
INTEGER, DIMENSION(:), POINTER :: unique_list_a
LOGICAL, INTENT(IN) :: use_virial
CHARACTER(LEN=*), PARAMETER :: routineN = 'allegro_energy_store_force_virial'
INTEGER :: atom_a, atom_b, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, ilast, ilist, &
ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, nedges, nloc_size, &
npairs, nunique
INTEGER :: atom_a, atom_b, atom_idx, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, &
ilast, ilist, ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, &
nedges, nloc_size, npairs, nunique
INTEGER(kind=int_8), ALLOCATABLE :: atom_types(:), temp_atom_types(:)
INTEGER(kind=int_8), ALLOCATABLE, DIMENSION(:, :) :: edge_index, t_edge_index, temp_edge_index
INTEGER, ALLOCATABLE, DIMENSION(:) :: work_list
@ -248,7 +251,8 @@ CONTAINS
REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: edge_cell_shifts, new_edge_cell_shifts, &
pos
REAL(kind=dp), DIMENSION(3) :: cell_v, cvi
REAL(kind=dp), DIMENSION(:, :), POINTER :: atomic_energy, forces
REAL(kind=dp), DIMENSION(:, :), POINTER :: atomic_energy, forces, virial
REAL(kind=dp), DIMENSION(:, :, :), POINTER :: virial3d
REAL(kind=sp) :: lattice_sp(3, 3)
REAL(kind=sp), ALLOCATABLE, DIMENSION(:, :) :: new_edge_cell_shifts_sp, pos_sp
REAL(kind=sp), DIMENSION(:, :), POINTER :: atomic_energy_sp, forces_sp
@ -259,22 +263,22 @@ CONTAINS
CALL timeset(routineN, handle)
NULLIFY (atomic_energy, forces, atomic_energy_sp, forces_sp)
NULLIFY (atomic_energy, forces, atomic_energy_sp, forces_sp, virial3d, virial)
n_atoms = SIZE(particle_set)
ALLOCATE (use_atom(n_atoms))
use_atom = .FALSE.
DO ikind = 1, SIZE(atomic_kind_set)
DO jkind = 1, SIZE(atomic_kind_set)
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= allegro_type) CYCLE
DO iat = 1, n_atoms
IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
END DO ! iat
END DO ! i
END DO ! jkind
DO jkind = 1, SIZE(atomic_kind_set)
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= allegro_type) CYCLE
DO iat = 1, n_atoms
IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
END DO ! iat
END DO ! i
END DO ! jkind
END DO ! ikind
n_atoms_use = COUNT(use_atom)
@ -394,17 +398,20 @@ CONTAINS
t_edge_index(:, :) = TRANSPOSE(temp_edge_index)
DEALLOCATE (temp_edge_index, edge_index)
lattice = cell%hmat/pot%set(1)%allegro%unit_cell_val
lattice_sp = REAL(lattice, kind=sp)
iat_use = 0
ALLOCATE (pos(3, n_atoms_use), atom_types(n_atoms_use))
DO iat = 1, n_atoms_use
IF (.NOT. use_atom(iat)) CYCLE
iat_use = iat_use + 1
atom_types(iat_use) = particle_set(iat)%atomic_kind%kind_number - 1
! Find index of the element based on its position in config.yaml file to have correct mapping
DO i = 1, SIZE(allegro%type_names_torch)
IF (particle_set(iat)%atomic_kind%element_symbol == allegro%type_names_torch(i)) THEN
atom_idx = i - 1
END IF
END DO
atom_types(iat_use) = atom_idx
pos(:, iat) = r_last_update_pbc(iat)%r(:)/allegro%unit_coords_val
END DO
@ -423,17 +430,22 @@ CONTAINS
CALL torch_dict_insert(inputs, "edge_cell_shift", new_edge_cell_shifts)
CALL torch_dict_insert(inputs, "cell", lattice)
END IF
CALL torch_dict_insert(inputs, "edge_index", t_edge_index)
CALL torch_dict_insert(inputs, "atom_types", atom_types)
CALL torch_dict_create(outputs)
CALL torch_model_eval(allegro_data%model, inputs, outputs)
pot_allegro = 0.0_dp
IF (allegro%do_allegro_sp) THEN
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy_sp)
CALL torch_dict_get(outputs, "forces", forces_sp)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
allegro_data%virial(:, :) = virial(:, :)*allegro%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
allegro_data%force(:, :) = REAL(forces_sp(:, :), kind=dp)*allegro%unit_forces_val
DO iat_use = 1, SIZE(unique_list_a)
i = unique_list_a(iat_use)
@ -443,6 +455,13 @@ CONTAINS
ELSE
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
CALL torch_dict_get(outputs, "forces", forces)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
allegro_data%virial(:, :) = virial(:, :)*allegro%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
allegro_data%force(:, :) = forces(:, :)*allegro%unit_forces_val
DO iat_use = 1, SIZE(unique_list_a)
i = unique_list_a(iat_use)
@ -473,15 +492,12 @@ CONTAINS
LOGICAL, INTENT(IN) :: use_virial
INTEGER :: iat, iat_use
REAL(KIND=dp), DIMENSION(3, 3) :: virial
TYPE(allegro_data_type), POINTER :: allegro_data
CALL fist_nonbond_env_get(fist_nonbond_env, allegro_data=allegro_data)
IF (use_virial) THEN
virial = 0.0_dp
pv_nonbond = pv_nonbond + virial
CPABORT("Stress tensor for Allegro not yet implemented")
pv_nonbond = pv_nonbond + allegro_data%virial
END IF
DO iat_use = 1, SIZE(allegro_data%use_indices)

65
src/manybody_nequip.F
View file

@ -182,7 +182,6 @@ CONTAINS
END IF
END SUBROUTINE destroy_nequip_arrays
! **************************************************************************************************
!> \brief ...
!> \param nonbonded ...
@ -196,13 +195,15 @@ CONTAINS
!> \param pot_nequip ...
!> \param fist_nonbond_env ...
!> \param para_env ...
!> \param use_virial ...
!> \par History
!> Implementation of the nequip potential - [gtocci] 2022
!> Index mapping of atoms from .xyz to Allegro config.yaml file - [mbilichenko] 2024
!> \author Gabriele Tocci - University of Zurich
! **************************************************************************************************
SUBROUTINE nequip_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, &
potparm, nequip, glob_loc_list_a, r_last_update_pbc, &
pot_nequip, fist_nonbond_env, para_env)
pot_nequip, fist_nonbond_env, para_env, use_virial)
TYPE(fist_neighbor_type), POINTER :: nonbonded
TYPE(particle_type), POINTER :: particle_set(:)
@ -215,12 +216,13 @@ CONTAINS
REAL(kind=dp) :: pot_nequip
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
TYPE(mp_para_env_type), OPTIONAL, POINTER :: para_env
LOGICAL, INTENT(IN) :: use_virial
CHARACTER(LEN=*), PARAMETER :: routineN = 'nequip_energy_store_force_virial'
INTEGER :: atom_a, atom_b, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, ilast, ilist, &
ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, nedges, nedges_tot, &
nloc_size, npairs, nunique
INTEGER :: atom_a, atom_b, atom_idx, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, &
ilast, ilist, ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, &
nedges, nedges_tot, nloc_size, npairs, nunique
INTEGER(kind=int_8), ALLOCATABLE :: atom_types(:)
INTEGER(kind=int_8), ALLOCATABLE, DIMENSION(:, :) :: edge_index, t_edge_index, temp_edge_index
INTEGER, ALLOCATABLE, DIMENSION(:) :: displ, displ_cell, edge_count, &
@ -231,7 +233,9 @@ CONTAINS
REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: edge_cell_shifts, pos, &
temp_edge_cell_shifts
REAL(kind=dp), DIMENSION(3) :: cell_v, cvi
REAL(kind=dp), DIMENSION(:, :), POINTER :: atomic_energy, forces, total_energy
REAL(kind=dp), DIMENSION(:, :), POINTER :: atomic_energy, forces, total_energy, &
virial
REAL(kind=dp), DIMENSION(:, :, :), POINTER :: virial3d
REAL(kind=sp) :: lattice_sp(3, 3)
REAL(kind=sp), ALLOCATABLE, DIMENSION(:, :) :: edge_cell_shifts_sp, pos_sp
REAL(kind=sp), DIMENSION(:, :), POINTER :: atomic_energy_sp, forces_sp, &
@ -243,22 +247,22 @@ CONTAINS
CALL timeset(routineN, handle)
NULLIFY (total_energy, atomic_energy, forces, total_energy_sp, atomic_energy_sp, forces_sp)
NULLIFY (total_energy, atomic_energy, forces, total_energy_sp, atomic_energy_sp, forces_sp, virial3d, virial)
n_atoms = SIZE(particle_set)
ALLOCATE (use_atom(n_atoms))
use_atom = .FALSE.
DO ikind = 1, SIZE(atomic_kind_set)
DO jkind = 1, SIZE(atomic_kind_set)
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= nequip_type) CYCLE
DO iat = 1, n_atoms
IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
END DO ! iat
END DO ! i
END DO ! jkind
DO jkind = 1, SIZE(atomic_kind_set)
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= nequip_type) CYCLE
DO iat = 1, n_atoms
IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
END DO ! iat
END DO ! i
END DO ! jkind
END DO ! ikind
n_atoms_use = COUNT(use_atom)
@ -408,7 +412,13 @@ CONTAINS
DO iat = 1, n_atoms_use
IF (.NOT. use_atom(iat)) CYCLE
iat_use = iat_use + 1
atom_types(iat_use) = particle_set(iat)%atomic_kind%kind_number - 1
! Find index of the element based on its position in config.yaml file to have correct mapping
DO i = 1, SIZE(nequip%type_names_torch)
IF (particle_set(iat)%atomic_kind%element_symbol == nequip%type_names_torch(i)) THEN
atom_idx = i - 1
END IF
END DO
atom_types(iat_use) = atom_idx
pos(:, iat) = r_last_update_pbc(iat)%r(:)/nequip%unit_coords_val
END DO
@ -437,6 +447,13 @@ CONTAINS
CALL torch_dict_get(outputs, "total_energy", total_energy_sp)
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy_sp)
CALL torch_dict_get(outputs, "forces", forces_sp)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
nequip_data%virial(:, :) = virial(:, :)*nequip%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
pot_nequip = REAL(total_energy_sp(1, 1), kind=dp)*nequip%unit_energy_val
nequip_data%force(:, :) = REAL(forces_sp(:, :), kind=dp)*nequip%unit_forces_val
DEALLOCATE (pos_sp, edge_cell_shifts_sp, total_energy_sp, atomic_energy_sp, forces_sp)
@ -444,6 +461,13 @@ CONTAINS
CALL torch_dict_get(outputs, "total_energy", total_energy)
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
CALL torch_dict_get(outputs, "forces", forces)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
nequip_data%virial(:, :) = virial(:, :)*nequip%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
pot_nequip = total_energy(1, 1)*nequip%unit_energy_val
nequip_data%force(:, :) = forces(:, :)*nequip%unit_forces_val
DEALLOCATE (pos, edge_cell_shifts, total_energy, atomic_energy, forces)
@ -477,15 +501,12 @@ CONTAINS
LOGICAL, INTENT(IN) :: use_virial
INTEGER :: iat, iat_use
REAL(KIND=dp), DIMENSION(3, 3) :: virial
TYPE(nequip_data_type), POINTER :: nequip_data
CALL fist_nonbond_env_get(fist_nonbond_env, nequip_data=nequip_data)
IF (use_virial) THEN
virial = 0.0_dp
pv_nonbond = pv_nonbond + virial
CPABORT("Stress tensor for NequIP not yet implemented")
pv_nonbond = pv_nonbond + nequip_data%virial
END IF
DO iat_use = 1, SIZE(nequip_data%use_indices)

8
src/manybody_potential.F
View file

@ -86,12 +86,13 @@ CONTAINS
!> \param pot_manybody ...
!> \param para_env ...
!> \param mm_section ...
!> \param use_virial ...
!> \par History
!> tlaino [2007] - New algorithm for tersoff potential
!> \author CJM, I-Feng W. Kuo, Teodoro Laino
! **************************************************************************************************
SUBROUTINE energy_manybody(fist_nonbond_env, atomic_kind_set, local_particles, &
particle_set, cell, pot_manybody, para_env, mm_section)
particle_set, cell, pot_manybody, para_env, mm_section, use_virial)
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
TYPE(atomic_kind_type), POINTER :: atomic_kind_set(:)
@ -101,6 +102,7 @@ CONTAINS
REAL(dp), INTENT(INOUT) :: pot_manybody
TYPE(mp_para_env_type), OPTIONAL, POINTER :: para_env
TYPE(section_vals_type), POINTER :: mm_section
LOGICAL, INTENT(IN) :: use_virial
CHARACTER(LEN=*), PARAMETER :: routineN = 'energy_manybody'
@ -205,7 +207,7 @@ CONTAINS
CALL setup_nequip_arrays(nonbonded, potparm, glob_loc_list, glob_cell_v, glob_loc_list_a, cell)
CALL nequip_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, potparm, &
nequip, glob_loc_list_a, r_last_update_pbc, pot_nequip, &
fist_nonbond_env, para_env)
fist_nonbond_env, para_env, use_virial)
pot_manybody = pot_manybody + pot_nequip
CALL destroy_nequip_arrays(glob_loc_list, glob_cell_v, glob_loc_list_a)
END IF
@ -216,7 +218,7 @@ CONTAINS
unique_list_a, cell)
CALL allegro_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, potparm, &
allegro, glob_loc_list_a, r_last_update_pbc, pot_allegro, &
fist_nonbond_env, unique_list_a)
fist_nonbond_env, unique_list_a, use_virial)
pot_manybody = pot_manybody + pot_allegro
CALL destroy_allegro_arrays(glob_loc_list, glob_cell_v, glob_loc_list_a, unique_list_a)
END IF

5
src/nequip_unittest.F
View file

@ -153,7 +153,7 @@ PROGRAM nequip_unittest
WRITE (*, *) "CUDA is available: ", torch_cuda_is_available()
filename = discover_file('NequIP/water.pth')
filename = discover_file('NequIP/water-deployed-neq060sp.pth')
WRITE (*, *) "Loading NequIP model from: "//TRIM(filename)
CALL torch_model_load(model, filename)
cutoff_str = torch_model_read_metadata(filename, "r_max")
@ -186,7 +186,8 @@ PROGRAM nequip_unittest
DO iatom = 1, natoms
WRITE (*, *) forces(:, iatom)*angstrom/evolt
END DO
CPASSERT(ABS(-14985.615_dp - REAL(total_energy(1, 1), kind=dp)) < 2e-3_dp)
CPASSERT(ABS(-14985.6299_dp - REAL(total_energy(1, 1), kind=dp)) < 2e-3_dp)
CALL torch_dict_release(inputs)
CALL torch_dict_release(outputs)

4
src/pair_potential_types.F
View file

@ -193,6 +193,7 @@ MODULE pair_potential_types
CHARACTER(LEN=default_path_length) :: nequip_file_name = 'NULL', nequip_version = 'NULL', &
unit_coords = 'NULL', unit_forces = 'NULL', &
unit_energy = 'NULL', unit_cell = 'NULL'
CHARACTER(LEN=100), DIMENSION(:), ALLOCATABLE :: type_names_torch
REAL(KIND=dp) :: rcutsq = 0.0_dp, unit_coords_val = 1.0_dp, &
unit_forces_val = 1.0_dp, unit_energy_val = 1.0_dp, &
unit_cell_val = 1.0_dp
@ -204,6 +205,9 @@ MODULE pair_potential_types
CHARACTER(LEN=default_path_length) :: allegro_file_name = 'NULL', unit_cell = 'NULL', &
nequip_version = 'NULL', unit_coords = 'NULL', &
unit_forces = 'NULL', unit_energy = 'NULL'
CHARACTER(LEN=100), DIMENSION(:), ALLOCATABLE :: type_names_torch
REAL(KIND=dp) :: rcutsq = 0.0_dp, unit_coords_val = 1.0_dp, &
unit_forces_val = 1.0_dp, unit_cell_val = 1.0_dp, &
unit_energy_val = 1.0_dp

4
tests/Fist/regtest-allegro/Allegro_si_MD.inp → tests/Fist/regtest-allegro/Si-dp.inp
View file

@ -34,7 +34,7 @@
&NONBONDED
&ALLEGRO
ATOMS Si
PARM_FILE_NAME Allegro/si-deployed.pth
PARM_FILE_NAME Allegro/si-deployed-neq060dp.pth
UNIT_COORDS angstrom
UNIT_ENERGY eV
UNIT_FORCES eV*angstrom^-1
@ -59,7 +59,7 @@
&TOPOLOGY
COORD_FILE_FORMAT XYZ
# coordinates must be ordered by atomic number
COORD_FILE_NAME si_input.xyz
COORD_FILE_NAME Si.xyz
# MULTIPLE_UNIT_CELL 4 4 4
&END TOPOLOGY
&END SUBSYS

0
tests/Fist/regtest-allegro/si_input.xyz → tests/Fist/regtest-allegro/Si.xyz
View file

5
tests/Fist/regtest-allegro/TEST_FILES
View file

@ -1,4 +1,5 @@
# Test of Allegro using libtorch https://pytorch.org/cppdocs/installing.html
Allegro_si_MD.inp 11 1.0E-6 -305.899211744308900
Allegro_wat_gra_film_double.inp 11 1.0E-12 -2073.467026976485613
Si-dp.inp 11 1.0E-12 -305.901401121683080
water-gra-film-sp.inp 11 1.0E-6 -2073.192819118499756
water-bulk-sp.inp 11 1.0E-6 -558.065488338470459
#EOF

53
tests/Fist/regtest-allegro/water-bulk-sp.inp Normal file
View file

@ -0,0 +1,53 @@
&GLOBAL
PRINT_LEVEL LOW
PROJECT water
RUN_TYPE MD
&END GLOBAL
&MOTION
&MD
ENSEMBLE NVT
STEPS 0
TEMPERATURE 300
TIMESTEP 0.5
&THERMOSTAT
&CSVR
TIMECON 10
&END CSVR
&END THERMOSTAT
&END MD
&END MOTION
&FORCE_EVAL
METHOD FIST
&MM
&FORCEFIELD
&NONBONDED
&ALLEGRO
ATOMS H O
PARM_FILE_NAME Allegro/gra-water-deployed-neq060sp.pth
UNIT_COORDS angstrom
UNIT_ENERGY Hartree
UNIT_FORCES eV*angstrom^-1
&END ALLEGRO
&END NONBONDED
&END FORCEFIELD
&POISSON
&EWALD
EWALD_TYPE none
&END EWALD
&END POISSON
&END MM
&SUBSYS
&CELL
A 9.8528 0.0 0.0
B 0.0 9.8528 0.0
C 0.0 0.0 9.8528
&END CELL
&TOPOLOGY
COORD_FILE_FORMAT XYZ
# coordinates must be ordered by atomic number
COORD_FILE_NAME ../regtest-nequip/water-bulk.xyz
&END TOPOLOGY
&END SUBSYS
&END FORCE_EVAL

5
tests/Fist/regtest-allegro/Allegro_wat_gra_film_double.inp → tests/Fist/regtest-allegro/water-gra-film-sp.inp
View file

@ -25,7 +25,7 @@
&NONBONDED
&ALLEGRO
ATOMS H C O
PARM_FILE_NAME Allegro/water-gra-film-double.pth
PARM_FILE_NAME Allegro/gra-water-deployed-neq060sp.pth
UNIT_COORDS angstrom
UNIT_ENERGY Hartree
UNIT_FORCES eV*angstrom^-1
@ -46,8 +46,7 @@
&END CELL
&TOPOLOGY
COORD_FILE_FORMAT XYZ
# coordinates must be ordered by atomic number
COORD_FILE_NAME ./wat_gra_film.xyz
COORD_FILE_NAME ./water-gra-film.xyz
&END TOPOLOGY
&END SUBSYS
&END FORCE_EVAL

0
tests/Fist/regtest-allegro/wat_gra_film.xyz → tests/Fist/regtest-allegro/water-gra-film.xyz
View file

5
tests/Fist/regtest-nequip/TEST_FILES
View file

@ -1,4 +1,5 @@
# Test of NequIP using libtorch https://pytorch.org/cppdocs/installing.html
NequIP_water.inp 11 1.0E-6 -17.194825205515503
NequIP_water_double.inp 11 1.0E-12 -550.638647597308932
water-sp.inp 11 1.0E-6 -17.192060707007283
water-bulk-dp.inp 11 1.0E-12 -550.631245735074231
water-bulk-dp.inp 31 1.0E-9 3.66494470310E+00
#EOF

12
tests/Fist/regtest-nequip/NequIP_water_double.inp → tests/Fist/regtest-nequip/water-bulk-dp.inp
View file

@ -19,12 +19,13 @@
&FORCE_EVAL
METHOD FIST
STRESS_TENSOR ANALYTICAL
&MM
&FORCEFIELD
&NONBONDED
&NEQUIP
ATOMS H O
PARM_FILE_NAME NequIP/water-double.pth
PARM_FILE_NAME NequIP/water-deployed-neq060dp.pth
UNIT_COORDS angstrom
UNIT_ENERGY eV
UNIT_FORCES eV*angstrom^-1
@ -40,17 +41,20 @@
&PRINT
&FORCES
&END FORCES
&STRESS_TENSOR
&EACH
MD 1
&END EACH
&END STRESS_TENSOR
&END PRINT
&SUBSYS
&CELL
ABC 9.8528 9.8528 9.8528
# MULTIPLE_UNIT_CELL 4 4 4
&END CELL
&TOPOLOGY
COORD_FILE_FORMAT XYZ
# MULTIPLE_UNIT_CELL 4 4 4
# coordinates must be ordered by atomic number
COORD_FILE_NAME ./H2O-32.xyz
COORD_FILE_NAME ./water-bulk.xyz
&END TOPOLOGY
&END SUBSYS
&END FORCE_EVAL

0
tests/Fist/regtest-nequip/H2O-32.xyz → tests/Fist/regtest-nequip/water-bulk.xyz
View file

2
tests/Fist/regtest-nequip/NequIP_water.inp → tests/Fist/regtest-nequip/water-sp.inp
View file

@ -12,7 +12,7 @@
&NONBONDED
&NEQUIP
ATOMS H O
PARM_FILE_NAME NequIP/water.pth
PARM_FILE_NAME NequIP/water-deployed-neq060sp.pth
UNIT_COORDS angstrom
UNIT_ENERGY eV
UNIT_FORCES eV*angstrom^-1