LCOV - code coverage report
Current view: top level - src - manybody_allegro.F (source / functions) Coverage Total Hit
Test: CP2K Regtests (git:42dac4a) Lines: 97.1 % 277 269
Test Date: 2025-07-25 12:55:17 Functions: 100.0 % 4 4

            Line data    Source code
       1              : !--------------------------------------------------------------------------------------------------!
       2              : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3              : !   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
       4              : !                                                                                                  !
       5              : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6              : !--------------------------------------------------------------------------------------------------!
       7              : 
       8              : ! **************************************************************************************************
       9              : !> \par History
      10              : !>      allegro implementation
      11              : !> \author Gabriele Tocci
      12              : ! **************************************************************************************************
      13              : MODULE manybody_allegro
      14              : 
      15              :    USE atomic_kind_types,               ONLY: atomic_kind_type
      16              :    USE cell_types,                      ONLY: cell_type
      17              :    USE fist_neighbor_list_types,        ONLY: fist_neighbor_type,&
      18              :                                               neighbor_kind_pairs_type
      19              :    USE fist_nonbond_env_types,          ONLY: allegro_data_type,&
      20              :                                               fist_nonbond_env_get,&
      21              :                                               fist_nonbond_env_set,&
      22              :                                               fist_nonbond_env_type,&
      23              :                                               pos_type
      24              :    USE kinds,                           ONLY: dp,&
      25              :                                               int_8,&
      26              :                                               sp
      27              :    USE message_passing,                 ONLY: mp_para_env_type
      28              :    USE pair_potential_types,            ONLY: allegro_pot_type,&
      29              :                                               allegro_type,&
      30              :                                               pair_potential_pp_type,&
      31              :                                               pair_potential_single_type
      32              :    USE particle_types,                  ONLY: particle_type
      33              :    USE torch_api,                       ONLY: &
      34              :         torch_dict_create, torch_dict_get, torch_dict_insert, torch_dict_release, torch_dict_type, &
      35              :         torch_model_forward, torch_model_freeze, torch_model_load, torch_tensor_data_ptr, &
      36              :         torch_tensor_from_array, torch_tensor_release, torch_tensor_type
      37              :    USE util,                            ONLY: sort
      38              : #include "./base/base_uses.f90"
      39              : 
      40              :    IMPLICIT NONE
      41              : 
      42              :    PRIVATE
      43              :    PUBLIC :: setup_allegro_arrays, destroy_allegro_arrays, &
      44              :              allegro_energy_store_force_virial, allegro_add_force_virial
      45              :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'manybody_allegro'
      46              : 
      47              : CONTAINS
      48              : 
      49              : ! **************************************************************************************************
      50              : !> \brief ...
      51              : !> \param nonbonded ...
      52              : !> \param potparm ...
      53              : !> \param glob_loc_list ...
      54              : !> \param glob_cell_v ...
      55              : !> \param glob_loc_list_a ...
      56              : !> \param unique_list_a ...
      57              : !> \param cell ...
      58              : !> \par History
      59              : !>      Implementation of the allegro potential - [gtocci] 2023
      60              : !> \author Gabriele Tocci - University of Zurich
      61              : ! **************************************************************************************************
      62            4 :    SUBROUTINE setup_allegro_arrays(nonbonded, potparm, glob_loc_list, glob_cell_v, glob_loc_list_a, &
      63              :                                    unique_list_a, cell)
      64              :       TYPE(fist_neighbor_type), POINTER                  :: nonbonded
      65              :       TYPE(pair_potential_pp_type), POINTER              :: potparm
      66              :       INTEGER, DIMENSION(:, :), POINTER                  :: glob_loc_list
      67              :       REAL(KIND=dp), DIMENSION(:, :), POINTER            :: glob_cell_v
      68              :       INTEGER, DIMENSION(:), POINTER                     :: glob_loc_list_a, unique_list_a
      69              :       TYPE(cell_type), POINTER                           :: cell
      70              : 
      71              :       CHARACTER(LEN=*), PARAMETER :: routineN = 'setup_allegro_arrays'
      72              : 
      73              :       INTEGER                                            :: handle, i, iend, igrp, ikind, ilist, &
      74              :                                                             ipair, istart, jkind, nkinds, nlocal, &
      75              :                                                             npairs, npairs_tot
      76            4 :       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: temp_unique_list_a, work_list, work_list2
      77            4 :       INTEGER, DIMENSION(:, :), POINTER                  :: list
      78            4 :       REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: rwork_list
      79              :       REAL(KIND=dp), DIMENSION(3)                        :: cell_v, cvi
      80              :       TYPE(neighbor_kind_pairs_type), POINTER            :: neighbor_kind_pair
      81              :       TYPE(pair_potential_single_type), POINTER          :: pot
      82              : 
      83            0 :       CPASSERT(.NOT. ASSOCIATED(glob_loc_list))
      84            4 :       CPASSERT(.NOT. ASSOCIATED(glob_loc_list_a))
      85            4 :       CPASSERT(.NOT. ASSOCIATED(unique_list_a))
      86            4 :       CPASSERT(.NOT. ASSOCIATED(glob_cell_v))
      87            4 :       CALL timeset(routineN, handle)
      88            4 :       npairs_tot = 0
      89            4 :       nkinds = SIZE(potparm%pot, 1)
      90          112 :       DO ilist = 1, nonbonded%nlists
      91          108 :          neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
      92          108 :          npairs = neighbor_kind_pair%npairs
      93          108 :          IF (npairs == 0) CYCLE
      94          258 :          Kind_Group_Loop1: DO igrp = 1, neighbor_kind_pair%ngrp_kind
      95          169 :             istart = neighbor_kind_pair%grp_kind_start(igrp)
      96          169 :             iend = neighbor_kind_pair%grp_kind_end(igrp)
      97          169 :             ikind = neighbor_kind_pair%ij_kind(1, igrp)
      98          169 :             jkind = neighbor_kind_pair%ij_kind(2, igrp)
      99          169 :             pot => potparm%pot(ikind, jkind)%pot
     100          169 :             npairs = iend - istart + 1
     101          169 :             IF (pot%no_mb) CYCLE
     102          446 :             DO i = 1, SIZE(pot%type)
     103          338 :                IF (pot%type(i) == allegro_type) npairs_tot = npairs_tot + npairs
     104              :             END DO
     105              :          END DO Kind_Group_Loop1
     106              :       END DO
     107           12 :       ALLOCATE (work_list(npairs_tot))
     108            8 :       ALLOCATE (work_list2(npairs_tot))
     109           12 :       ALLOCATE (glob_loc_list(2, npairs_tot))
     110           12 :       ALLOCATE (glob_cell_v(3, npairs_tot))
     111              :       ! Fill arrays with data
     112            4 :       npairs_tot = 0
     113          112 :       DO ilist = 1, nonbonded%nlists
     114          108 :          neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
     115          108 :          npairs = neighbor_kind_pair%npairs
     116          108 :          IF (npairs == 0) CYCLE
     117          258 :          Kind_Group_Loop2: DO igrp = 1, neighbor_kind_pair%ngrp_kind
     118          169 :             istart = neighbor_kind_pair%grp_kind_start(igrp)
     119          169 :             iend = neighbor_kind_pair%grp_kind_end(igrp)
     120          169 :             ikind = neighbor_kind_pair%ij_kind(1, igrp)
     121          169 :             jkind = neighbor_kind_pair%ij_kind(2, igrp)
     122          169 :             list => neighbor_kind_pair%list
     123          676 :             cvi = neighbor_kind_pair%cell_vector
     124          169 :             pot => potparm%pot(ikind, jkind)%pot
     125          169 :             npairs = iend - istart + 1
     126          169 :             IF (pot%no_mb) CYCLE
     127         2197 :             cell_v = MATMUL(cell%hmat, cvi)
     128          446 :             DO i = 1, SIZE(pot%type)
     129              :                ! ALLEGRO
     130          338 :                IF (pot%type(i) == allegro_type) THEN
     131        10533 :                   DO ipair = 1, npairs
     132        62184 :                      glob_loc_list(:, npairs_tot + ipair) = list(:, istart - 1 + ipair)
     133        41625 :                      glob_cell_v(1:3, npairs_tot + ipair) = cell_v(1:3)
     134              :                   END DO
     135          169 :                   npairs_tot = npairs_tot + npairs
     136              :                END IF
     137              :             END DO
     138              :          END DO Kind_Group_Loop2
     139              :       END DO
     140              :       ! Order the arrays w.r.t. the first index of glob_loc_list
     141            4 :       CALL sort(glob_loc_list(1, :), npairs_tot, work_list)
     142        10368 :       DO ipair = 1, npairs_tot
     143        10368 :          work_list2(ipair) = glob_loc_list(2, work_list(ipair))
     144              :       END DO
     145        10368 :       glob_loc_list(2, :) = work_list2
     146            4 :       DEALLOCATE (work_list2)
     147           12 :       ALLOCATE (rwork_list(3, npairs_tot))
     148        10368 :       DO ipair = 1, npairs_tot
     149        41460 :          rwork_list(:, ipair) = glob_cell_v(:, work_list(ipair))
     150              :       END DO
     151        41460 :       glob_cell_v = rwork_list
     152            4 :       DEALLOCATE (rwork_list)
     153            4 :       DEALLOCATE (work_list)
     154           12 :       ALLOCATE (glob_loc_list_a(npairs_tot))
     155        20736 :       glob_loc_list_a = glob_loc_list(1, :)
     156            8 :       ALLOCATE (temp_unique_list_a(npairs_tot))
     157            4 :       nlocal = 1
     158            4 :       temp_unique_list_a(1) = glob_loc_list_a(1)
     159        10364 :       DO ipair = 2, npairs_tot
     160        10364 :          IF (glob_loc_list_a(ipair - 1) /= glob_loc_list_a(ipair)) THEN
     161          156 :             nlocal = nlocal + 1
     162          156 :             temp_unique_list_a(nlocal) = glob_loc_list_a(ipair)
     163              :          END IF
     164              :       END DO
     165           12 :       ALLOCATE (unique_list_a(nlocal))
     166          164 :       unique_list_a(:) = temp_unique_list_a(:nlocal)
     167            4 :       DEALLOCATE (temp_unique_list_a)
     168            4 :       CALL timestop(handle)
     169            8 :    END SUBROUTINE setup_allegro_arrays
     170              : 
     171              : ! **************************************************************************************************
     172              : !> \brief ...
     173              : !> \param glob_loc_list ...
     174              : !> \param glob_cell_v ...
     175              : !> \param glob_loc_list_a ...
     176              : !> \param unique_list_a ...
     177              : !> \par History
     178              : !>      Implementation of the allegro potential - [gtocci] 2023
     179              : !> \author Gabriele Tocci - University of Zurich
     180              : ! **************************************************************************************************
     181            4 :    SUBROUTINE destroy_allegro_arrays(glob_loc_list, glob_cell_v, glob_loc_list_a, unique_list_a)
     182              :       INTEGER, DIMENSION(:, :), POINTER                  :: glob_loc_list
     183              :       REAL(KIND=dp), DIMENSION(:, :), POINTER            :: glob_cell_v
     184              :       INTEGER, DIMENSION(:), POINTER                     :: glob_loc_list_a, unique_list_a
     185              : 
     186            4 :       IF (ASSOCIATED(glob_loc_list)) THEN
     187            4 :          DEALLOCATE (glob_loc_list)
     188              :       END IF
     189            4 :       IF (ASSOCIATED(glob_loc_list_a)) THEN
     190            4 :          DEALLOCATE (glob_loc_list_a)
     191              :       END IF
     192            4 :       IF (ASSOCIATED(glob_cell_v)) THEN
     193            4 :          DEALLOCATE (glob_cell_v)
     194              :       END IF
     195            4 :       IF (ASSOCIATED(unique_list_a)) THEN
     196            4 :          DEALLOCATE (unique_list_a)
     197              :       END IF
     198              : 
     199            4 :    END SUBROUTINE destroy_allegro_arrays
     200              : 
     201              : ! **************************************************************************************************
     202              : !> \brief ...
     203              : !> \param nonbonded ...
     204              : !> \param particle_set ...
     205              : !> \param cell ...
     206              : !> \param atomic_kind_set ...
     207              : !> \param potparm ...
     208              : !> \param allegro ...
     209              : !> \param glob_loc_list_a ...
     210              : !> \param r_last_update_pbc ...
     211              : !> \param pot_allegro ...
     212              : !> \param fist_nonbond_env ...
     213              : !> \param unique_list_a ...
     214              : !> \param para_env ...
     215              : !> \param use_virial ...
     216              : !> \par History
     217              : !>      Implementation of the allegro potential - [gtocci] 2023
     218              : !>      Index mapping of atoms from .xyz to Allegro config.yaml file - [mbilichenko] 2024
     219              : !> \author Gabriele Tocci - University of Zurich
     220              : ! **************************************************************************************************
     221            4 :    SUBROUTINE allegro_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, &
     222              :                                                 potparm, allegro, glob_loc_list_a, r_last_update_pbc, &
     223              :                                                 pot_allegro, fist_nonbond_env, unique_list_a, para_env, use_virial)
     224              : 
     225              :       TYPE(fist_neighbor_type), POINTER                  :: nonbonded
     226              :       TYPE(particle_type), POINTER                       :: particle_set(:)
     227              :       TYPE(cell_type), POINTER                           :: cell
     228              :       TYPE(atomic_kind_type), POINTER                    :: atomic_kind_set(:)
     229              :       TYPE(pair_potential_pp_type), POINTER              :: potparm
     230              :       TYPE(allegro_pot_type), POINTER                    :: allegro
     231              :       INTEGER, DIMENSION(:), POINTER                     :: glob_loc_list_a
     232              :       TYPE(pos_type), DIMENSION(:), POINTER              :: r_last_update_pbc
     233              :       REAL(kind=dp)                                      :: pot_allegro
     234              :       TYPE(fist_nonbond_env_type), POINTER               :: fist_nonbond_env
     235              :       INTEGER, DIMENSION(:), POINTER                     :: unique_list_a
     236              :       TYPE(mp_para_env_type), POINTER                    :: para_env
     237              :       LOGICAL, INTENT(IN)                                :: use_virial
     238              : 
     239              :       CHARACTER(LEN=*), PARAMETER :: routineN = 'allegro_energy_store_force_virial'
     240              : 
     241              :       INTEGER :: atom_a, atom_b, atom_idx, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, &
     242              :          ilast, ilist, ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, &
     243              :          nedges, nloc_size, npairs, nunique
     244            4 :       INTEGER(kind=int_8), ALLOCATABLE                   :: atom_types(:), temp_atom_types(:)
     245            4 :       INTEGER(kind=int_8), ALLOCATABLE, DIMENSION(:, :)  :: edge_index, t_edge_index, temp_edge_index
     246            4 :       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: work_list
     247            4 :       INTEGER, DIMENSION(:, :), POINTER                  :: list, sort_list
     248            4 :       LOGICAL, ALLOCATABLE                               :: use_atom(:)
     249              :       REAL(kind=dp)                                      :: drij, rab2_max, rij(3)
     250            4 :       REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: edge_cell_shifts, lattice, &
     251            4 :                                                             new_edge_cell_shifts, pos
     252              :       REAL(kind=dp), DIMENSION(3)                        :: cell_v, cvi
     253            4 :       REAL(kind=dp), DIMENSION(:, :), POINTER            :: atomic_energy, forces, virial
     254            4 :       REAL(kind=dp), DIMENSION(:, :, :), POINTER         :: virial3d
     255            4 :       REAL(kind=sp), ALLOCATABLE, DIMENSION(:, :)        :: lattice_sp, new_edge_cell_shifts_sp, &
     256            4 :                                                             pos_sp
     257            4 :       REAL(kind=sp), DIMENSION(:, :), POINTER            :: atomic_energy_sp, forces_sp
     258              :       TYPE(allegro_data_type), POINTER                   :: allegro_data
     259              :       TYPE(neighbor_kind_pairs_type), POINTER            :: neighbor_kind_pair
     260              :       TYPE(pair_potential_single_type), POINTER          :: pot
     261              :       TYPE(torch_dict_type)                              :: inputs, outputs
     262              :       TYPE(torch_tensor_type) :: atom_types_tensor, atomic_energy_tensor, forces_tensor, &
     263              :          lattice_tensor, new_edge_cell_shifts_tensor, pos_tensor, t_edge_index_tensor, &
     264              :          virial_tensor
     265              : 
     266            4 :       CALL timeset(routineN, handle)
     267              : 
     268            4 :       NULLIFY (atomic_energy, forces, atomic_energy_sp, forces_sp, virial3d, virial)
     269            4 :       n_atoms = SIZE(particle_set)
     270           12 :       ALLOCATE (use_atom(n_atoms))
     271          324 :       use_atom = .FALSE.
     272              : 
     273           10 :       DO ikind = 1, SIZE(atomic_kind_set)
     274           20 :          DO jkind = 1, SIZE(atomic_kind_set)
     275           10 :             pot => potparm%pot(ikind, jkind)%pot
     276           26 :             DO i = 1, SIZE(pot%type)
     277           10 :                IF (pot%type(i) /= allegro_type) CYCLE
     278          916 :                DO iat = 1, n_atoms
     279          896 :                   IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
     280          714 :                       particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
     281              :                END DO ! iat
     282              :             END DO ! i
     283              :          END DO ! jkind
     284              :       END DO ! ikind
     285          324 :       n_atoms_use = COUNT(use_atom)
     286              : 
     287              :       ! get allegro_data to save force, virial info and to load model
     288            4 :       CALL fist_nonbond_env_get(fist_nonbond_env, allegro_data=allegro_data)
     289            4 :       IF (.NOT. ASSOCIATED(allegro_data)) THEN
     290           52 :          ALLOCATE (allegro_data)
     291            4 :          CALL fist_nonbond_env_set(fist_nonbond_env, allegro_data=allegro_data)
     292            4 :          NULLIFY (allegro_data%use_indices, allegro_data%force)
     293            4 :          CALL torch_model_load(allegro_data%model, pot%set(1)%allegro%allegro_file_name)
     294            4 :          CALL torch_model_freeze(allegro_data%model)
     295              :       END IF
     296            4 :       IF (ASSOCIATED(allegro_data%force)) THEN
     297            0 :          IF (SIZE(allegro_data%force, 2) /= n_atoms_use) THEN
     298            0 :             DEALLOCATE (allegro_data%force, allegro_data%use_indices)
     299              :          END IF
     300              :       END IF
     301            4 :       IF (.NOT. ASSOCIATED(allegro_data%force)) THEN
     302           12 :          ALLOCATE (allegro_data%force(3, n_atoms_use))
     303           12 :          ALLOCATE (allegro_data%use_indices(n_atoms_use))
     304              :       END IF
     305              : 
     306              :       iat_use = 0
     307          324 :       DO iat = 1, n_atoms_use
     308          324 :          IF (use_atom(iat)) THEN
     309          320 :             iat_use = iat_use + 1
     310          320 :             allegro_data%use_indices(iat_use) = iat
     311              :          END IF
     312              :       END DO
     313              : 
     314            4 :       nedges = 0
     315              : 
     316           12 :       ALLOCATE (edge_index(2, SIZE(glob_loc_list_a)))
     317           12 :       ALLOCATE (edge_cell_shifts(3, SIZE(glob_loc_list_a)))
     318           12 :       ALLOCATE (temp_atom_types(SIZE(glob_loc_list_a)))
     319              : 
     320          112 :       DO ilist = 1, nonbonded%nlists
     321          108 :          neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
     322          108 :          npairs = neighbor_kind_pair%npairs
     323          108 :          IF (npairs == 0) CYCLE
     324          258 :          Kind_Group_Loop_Allegro: DO igrp = 1, neighbor_kind_pair%ngrp_kind
     325          169 :             istart = neighbor_kind_pair%grp_kind_start(igrp)
     326          169 :             iend = neighbor_kind_pair%grp_kind_end(igrp)
     327          169 :             ikind = neighbor_kind_pair%ij_kind(1, igrp)
     328          169 :             jkind = neighbor_kind_pair%ij_kind(2, igrp)
     329          169 :             list => neighbor_kind_pair%list
     330          676 :             cvi = neighbor_kind_pair%cell_vector
     331          169 :             pot => potparm%pot(ikind, jkind)%pot
     332          446 :             DO i = 1, SIZE(pot%type)
     333          169 :                IF (pot%type(i) /= allegro_type) CYCLE
     334          169 :                rab2_max = pot%set(i)%allegro%rcutsq
     335         2197 :                cell_v = MATMUL(cell%hmat, cvi)
     336          169 :                pot => potparm%pot(ikind, jkind)%pot
     337          169 :                allegro => pot%set(i)%allegro
     338          169 :                npairs = iend - istart + 1
     339          338 :                IF (npairs /= 0) THEN
     340          845 :                   ALLOCATE (sort_list(2, npairs), work_list(npairs))
     341        62353 :                   sort_list = list(:, istart:iend)
     342              :                   ! Sort the list of neighbors, this increases the efficiency for single
     343              :                   ! potential contributions
     344          169 :                   CALL sort(sort_list(1, :), npairs, work_list)
     345        10533 :                   DO ipair = 1, npairs
     346        10533 :                      work_list(ipair) = sort_list(2, work_list(ipair))
     347              :                   END DO
     348        10533 :                   sort_list(2, :) = work_list
     349              :                   ! find number of unique elements of array index 1
     350              :                   nunique = 1
     351        10364 :                   DO ipair = 1, npairs - 1
     352        10364 :                      IF (sort_list(1, ipair + 1) /= sort_list(1, ipair)) nunique = nunique + 1
     353              :                   END DO
     354          169 :                   ipair = 1
     355          169 :                   junique = sort_list(1, ipair)
     356          169 :                   ifirst = 1
     357         1538 :                   DO iunique = 1, nunique
     358         1369 :                      atom_a = junique
     359         1369 :                      IF (glob_loc_list_a(ifirst) > atom_a) CYCLE
     360       360440 :                      DO mpair = ifirst, SIZE(glob_loc_list_a)
     361       360440 :                         IF (glob_loc_list_a(mpair) == atom_a) EXIT
     362              :                      END DO
     363       106529 :                      ifirst = mpair
     364       106529 :                      DO mpair = ifirst, SIZE(glob_loc_list_a)
     365       106529 :                         IF (glob_loc_list_a(mpair) /= atom_a) EXIT
     366              :                      END DO
     367         1369 :                      ilast = mpair - 1
     368         1369 :                      nloc_size = 0
     369         1369 :                      IF (ifirst /= 0) nloc_size = ilast - ifirst + 1
     370        11733 :                      DO WHILE (ipair <= npairs)
     371        11564 :                         IF (sort_list(1, ipair) /= junique) EXIT
     372        10364 :                         atom_b = sort_list(2, ipair)
     373        41456 :                         rij(:) = r_last_update_pbc(atom_b)%r(:) - r_last_update_pbc(atom_a)%r(:) + cell_v
     374        41456 :                         drij = DOT_PRODUCT(rij, rij)
     375        10364 :                         ipair = ipair + 1
     376        11733 :                         IF (drij <= rab2_max) THEN
     377         5998 :                            nedges = nedges + 1
     378        17994 :                            edge_index(:, nedges) = [atom_a - 1, atom_b - 1]
     379        23992 :                            edge_cell_shifts(:, nedges) = cvi
     380              :                         END IF
     381              :                      END DO
     382         1369 :                      ifirst = ilast + 1
     383         1538 :                      IF (ipair <= npairs) junique = sort_list(1, ipair)
     384              :                   END DO
     385          169 :                   DEALLOCATE (sort_list, work_list)
     386              :                END IF
     387              :             END DO
     388              :          END DO Kind_Group_Loop_Allegro
     389              :       END DO
     390              : 
     391            4 :       allegro => pot%set(1)%allegro
     392              : 
     393           12 :       ALLOCATE (temp_edge_index(2, nedges))
     394        17998 :       temp_edge_index(:, :) = edge_index(:, :nedges)
     395           12 :       ALLOCATE (new_edge_cell_shifts(3, nedges))
     396        23996 :       new_edge_cell_shifts(:, :) = edge_cell_shifts(:, :nedges)
     397            4 :       DEALLOCATE (edge_cell_shifts)
     398              : 
     399            8 :       ALLOCATE (t_edge_index(nedges, 2))
     400              : 
     401        12008 :       t_edge_index(:, :) = TRANSPOSE(temp_edge_index)
     402            4 :       DEALLOCATE (temp_edge_index, edge_index)
     403            4 :       ALLOCATE (lattice(3, 3), lattice_sp(3, 3))
     404           52 :       lattice(:, :) = cell%hmat/pot%set(1)%allegro%unit_cell_val
     405           52 :       lattice_sp(:, :) = REAL(lattice, kind=sp)
     406            4 :       iat_use = 0
     407           20 :       ALLOCATE (pos(3, n_atoms_use), atom_types(n_atoms_use))
     408          324 :       DO iat = 1, n_atoms_use
     409          320 :          IF (.NOT. use_atom(iat)) CYCLE
     410          320 :          iat_use = iat_use + 1
     411              :          ! Find index of the element based on its position in config.yaml file to have correct mapping
     412         1024 :          DO i = 1, SIZE(allegro%type_names_torch)
     413         1024 :             IF (particle_set(iat)%atomic_kind%element_symbol == allegro%type_names_torch(i)) THEN
     414          320 :                atom_idx = i - 1
     415              :             END IF
     416              :          END DO
     417          320 :          atom_types(iat_use) = atom_idx
     418         1284 :          pos(:, iat) = r_last_update_pbc(iat)%r(:)/allegro%unit_coords_val
     419              :       END DO
     420              : 
     421            4 :       CALL torch_dict_create(inputs)
     422              : 
     423            4 :       IF (allegro%do_allegro_sp) THEN
     424           10 :          ALLOCATE (new_edge_cell_shifts_sp(3, nedges), pos_sp(3, n_atoms_use))
     425        19898 :          new_edge_cell_shifts_sp(:, :) = REAL(new_edge_cell_shifts(:, :), kind=sp)
     426          770 :          pos_sp(:, :) = REAL(pos(:, :), kind=sp)
     427            2 :          DEALLOCATE (pos, new_edge_cell_shifts)
     428            2 :          CALL torch_tensor_from_array(pos_tensor, pos_sp)
     429            2 :          CALL torch_tensor_from_array(new_edge_cell_shifts_tensor, new_edge_cell_shifts_sp)
     430            2 :          CALL torch_tensor_from_array(lattice_tensor, lattice_sp)
     431              :       ELSE
     432            2 :          CALL torch_tensor_from_array(pos_tensor, pos)
     433            2 :          CALL torch_tensor_from_array(new_edge_cell_shifts_tensor, new_edge_cell_shifts)
     434            2 :          CALL torch_tensor_from_array(lattice_tensor, lattice)
     435              :       END IF
     436              : 
     437            4 :       CALL torch_dict_insert(inputs, "pos", pos_tensor)
     438            4 :       CALL torch_dict_insert(inputs, "edge_cell_shift", new_edge_cell_shifts_tensor)
     439            4 :       CALL torch_dict_insert(inputs, "cell", lattice_tensor)
     440            4 :       CALL torch_tensor_release(pos_tensor)
     441            4 :       CALL torch_tensor_release(new_edge_cell_shifts_tensor)
     442            4 :       CALL torch_tensor_release(lattice_tensor)
     443              : 
     444            4 :       CALL torch_tensor_from_array(t_edge_index_tensor, t_edge_index)
     445            4 :       CALL torch_dict_insert(inputs, "edge_index", t_edge_index_tensor)
     446            4 :       CALL torch_tensor_release(t_edge_index_tensor)
     447              : 
     448            4 :       CALL torch_tensor_from_array(atom_types_tensor, atom_types)
     449            4 :       CALL torch_dict_insert(inputs, "atom_types", atom_types_tensor)
     450            4 :       CALL torch_tensor_release(atom_types_tensor)
     451              : 
     452            4 :       CALL torch_dict_create(outputs)
     453            4 :       CALL torch_model_forward(allegro_data%model, inputs, outputs)
     454            4 :       pot_allegro = 0.0_dp
     455              : 
     456            4 :       CALL torch_dict_get(outputs, "atomic_energy", atomic_energy_tensor)
     457            4 :       CALL torch_dict_get(outputs, "forces", forces_tensor)
     458            4 :       IF (allegro%do_allegro_sp) THEN
     459            2 :          CALL torch_tensor_data_ptr(atomic_energy_tensor, atomic_energy_sp)
     460            2 :          CALL torch_tensor_data_ptr(forces_tensor, forces_sp)
     461          770 :          allegro_data%force(:, :) = REAL(forces_sp(:, :), kind=dp)*allegro%unit_forces_val
     462           98 :          DO iat_use = 1, SIZE(unique_list_a)
     463           96 :             i = unique_list_a(iat_use)
     464           98 :             pot_allegro = pot_allegro + REAL(atomic_energy_sp(1, i), kind=dp)*allegro%unit_energy_val
     465              :          END DO
     466            2 :          DEALLOCATE (new_edge_cell_shifts_sp, pos_sp)
     467              :       ELSE
     468            2 :          CALL torch_tensor_data_ptr(atomic_energy_tensor, atomic_energy)
     469            2 :          CALL torch_tensor_data_ptr(forces_tensor, forces)
     470              : 
     471         1026 :          allegro_data%force(:, :) = forces(:, :)*allegro%unit_forces_val
     472           66 :          DO iat_use = 1, SIZE(unique_list_a)
     473           64 :             i = unique_list_a(iat_use)
     474           66 :             pot_allegro = pot_allegro + atomic_energy(1, i)*allegro%unit_energy_val
     475              :          END DO
     476            2 :          DEALLOCATE (pos, new_edge_cell_shifts)
     477              :       END IF
     478            4 :       CALL torch_tensor_release(atomic_energy_tensor)
     479            4 :       CALL torch_tensor_release(forces_tensor)
     480              : 
     481            4 :       IF (use_virial) THEN
     482            0 :          CALL torch_dict_get(outputs, "virial", virial_tensor)
     483            0 :          CALL torch_tensor_data_ptr(virial_tensor, virial3d)
     484            0 :          allegro_data%virial(:, :) = RESHAPE(virial3d, (/3, 3/))*allegro%unit_energy_val
     485            0 :          CALL torch_tensor_release(virial_tensor)
     486              :       END IF
     487              : 
     488            4 :       CALL torch_dict_release(inputs)
     489            4 :       CALL torch_dict_release(outputs)
     490              : 
     491            4 :       DEALLOCATE (t_edge_index, atom_types)
     492              : 
     493            4 :       IF (use_virial) allegro_data%virial(:, :) = allegro_data%virial/REAL(para_env%num_pe, dp)
     494            4 :       CALL timestop(handle)
     495            8 :    END SUBROUTINE allegro_energy_store_force_virial
     496              : 
     497              : ! **************************************************************************************************
     498              : !> \brief ...
     499              : !> \param fist_nonbond_env ...
     500              : !> \param f_nonbond ...
     501              : !> \param pv_nonbond ...
     502              : !> \param use_virial ...
     503              : ! **************************************************************************************************
     504            4 :    SUBROUTINE allegro_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
     505              : 
     506              :       TYPE(fist_nonbond_env_type), POINTER               :: fist_nonbond_env
     507              :       REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT)      :: f_nonbond, pv_nonbond
     508              :       LOGICAL, INTENT(IN)                                :: use_virial
     509              : 
     510              :       INTEGER                                            :: iat, iat_use
     511              :       TYPE(allegro_data_type), POINTER                   :: allegro_data
     512              : 
     513            4 :       CALL fist_nonbond_env_get(fist_nonbond_env, allegro_data=allegro_data)
     514              : 
     515            4 :       IF (use_virial) THEN
     516            0 :          pv_nonbond = pv_nonbond + allegro_data%virial
     517              :       END IF
     518              : 
     519          324 :       DO iat_use = 1, SIZE(allegro_data%use_indices)
     520          320 :          iat = allegro_data%use_indices(iat_use)
     521          320 :          CPASSERT(iat >= 1 .AND. iat <= SIZE(f_nonbond, 2))
     522         1284 :          f_nonbond(1:3, iat) = f_nonbond(1:3, iat) + allegro_data%force(1:3, iat_use)
     523              :       END DO
     524              : 
     525            4 :    END SUBROUTINE allegro_add_force_virial
     526              : END MODULE manybody_allegro
     527              : 
        

Generated by: LCOV version 2.0-1