Line data    Source code

       1             : !--------------------------------------------------------------------------------------------------!
       2             : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3             : !   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
       4             : !                                                                                                  !
       5             : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6             : !--------------------------------------------------------------------------------------------------!
       7             : 
       8             : MODULE manybody_deepmd
       9             : 
      10             :    USE atomic_kind_types,               ONLY: atomic_kind_type
      11             :    USE bibliography,                    ONLY: Wang2018,&
      12             :                                               Zeng2023,&
      13             :                                               cite_reference
      14             :    USE cell_types,                      ONLY: cell_type
      15             :    USE cp_log_handling,                 ONLY: cp_logger_get_default_io_unit
      16             :    USE deepmd_wrapper,                  ONLY: deepmd_model_compute,&
      17             :                                               deepmd_model_load
      18             :    USE fist_nonbond_env_types,          ONLY: deepmd_data_type,&
      19             :                                               fist_nonbond_env_get,&
      20             :                                               fist_nonbond_env_set,&
      21             :                                               fist_nonbond_env_type
      22             :    USE kinds,                           ONLY: dp
      23             :    USE message_passing,                 ONLY: mp_para_env_type
      24             :    USE pair_potential_types,            ONLY: deepmd_type,&
      25             :                                               pair_potential_pp_type,&
      26             :                                               pair_potential_single_type
      27             :    USE particle_types,                  ONLY: particle_type
      28             :    USE physcon,                         ONLY: angstrom,&
      29             :                                               evolt
      30             : #include "./base/base_uses.f90"
      31             : 
      32             :    IMPLICIT NONE
      33             : 
      34             :    PRIVATE
      35             :    PUBLIC deepmd_energy_store_force_virial, deepmd_add_force_virial
      36             : 
      37             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'manybody_deepmd'
      38             : 
      39             : CONTAINS
      40             : 
      41             : ! **************************************************************************************************
      42             : !> \brief ...
      43             : !> \param particle_set ...
      44             : !> \param cell ...
      45             : !> \param atomic_kind_set ...
      46             : !> \param potparm ...
      47             : !> \param fist_nonbond_env ...
      48             : !> \param pot_deepmd ...
      49             : !> \param para_env ...
      50             : ! **************************************************************************************************
      51           2 :    SUBROUTINE deepmd_energy_store_force_virial(particle_set, cell, atomic_kind_set, potparm, fist_nonbond_env, &
      52             :                                                pot_deepmd, para_env)
      53             :       TYPE(particle_type), POINTER                       :: particle_set(:)
      54             :       TYPE(cell_type), POINTER                           :: cell
      55             :       TYPE(atomic_kind_type), POINTER                    :: atomic_kind_set(:)
      56             :       TYPE(pair_potential_pp_type), POINTER              :: potparm
      57             :       TYPE(fist_nonbond_env_type), POINTER               :: fist_nonbond_env
      58             :       REAL(kind=dp)                                      :: pot_deepmd
      59             :       TYPE(mp_para_env_type), OPTIONAL, POINTER          :: para_env
      60             : 
      61             :       CHARACTER(LEN=*), PARAMETER :: routineN = 'deepmd_energy_store_force_virial'
      62             : 
      63             :       INTEGER                                            :: dpmd_natoms, handle, i, iat, iat_use, &
      64             :                                                             ikind, jkind, n_atoms, n_atoms_use, &
      65             :                                                             output_unit
      66           2 :       INTEGER, ALLOCATABLE                               :: dpmd_atype(:), use_atom_type(:)
      67           2 :       LOGICAL, ALLOCATABLE                               :: use_atom(:)
      68             :       REAL(kind=dp)                                      :: lattice(3, 3)
      69           2 :       REAL(kind=dp), ALLOCATABLE                         :: dpmd_atomic_energy(:), &
      70           2 :                                                             dpmd_atomic_virial(:), dpmd_cell(:), &
      71           2 :                                                             dpmd_coord(:), dpmd_force(:), &
      72           2 :                                                             dpmd_virial(:)
      73             :       TYPE(deepmd_data_type), POINTER                    :: deepmd_data
      74             :       TYPE(pair_potential_single_type), POINTER          :: pot
      75             : 
      76           2 :       CALL timeset(routineN, handle)
      77           2 :       n_atoms = SIZE(particle_set)
      78           6 :       ALLOCATE (use_atom(n_atoms))
      79           4 :       ALLOCATE (use_atom_type(n_atoms))
      80          18 :       use_atom = .FALSE.
      81          18 :       use_atom_type = 0
      82             : 
      83           4 :       DO ikind = 1, SIZE(atomic_kind_set)
      84           6 :          DO jkind = 1, SIZE(atomic_kind_set)
      85           2 :             pot => potparm%pot(ikind, jkind)%pot
      86             :             ! ensure that each atom is only used once
      87           2 :             IF (ikind /= jkind) CYCLE
      88           6 :             DO i = 1, SIZE(pot%type)
      89           2 :                IF (pot%type(i) /= deepmd_type) CYCLE
      90          20 :                DO iat = 1, n_atoms
      91          16 :                   IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
      92           2 :                       particle_set(iat)%atomic_kind%kind_number == jkind) THEN
      93          16 :                      use_atom(iat) = .TRUE.
      94          16 :                      use_atom_type(iat) = pot%set(i)%deepmd%atom_deepmd_type
      95             :                   END IF
      96             :                END DO ! iat
      97             :             END DO ! i
      98             :          END DO ! jkind
      99             :       END DO ! ikind
     100             : 
     101          18 :       n_atoms_use = COUNT(use_atom)
     102           2 :       dpmd_natoms = n_atoms_use
     103             :       ALLOCATE (dpmd_cell(9), dpmd_atype(dpmd_natoms), dpmd_coord(dpmd_natoms*3), &
     104             :                 dpmd_force(dpmd_natoms*3), dpmd_virial(9), &
     105          20 :                 dpmd_atomic_energy(dpmd_natoms), dpmd_atomic_virial(dpmd_natoms*9))
     106             : 
     107          18 :       iat_use = 0
     108          18 :       DO iat = 1, n_atoms
     109          16 :          IF (.NOT. use_atom(iat)) CYCLE
     110          16 :          iat_use = iat_use + 1
     111          64 :          dpmd_coord((iat_use - 1)*3 + 1:(iat_use - 1)*3 + 3) = particle_set(iat)%r*angstrom
     112          18 :          dpmd_atype(iat_use) = use_atom_type(iat)
     113             :       END DO
     114           2 :       IF (iat_use > 0) THEN
     115           2 :          CALL cite_reference(Wang2018)
     116           2 :          CALL cite_reference(Zeng2023)
     117             :       END IF
     118           2 :       output_unit = cp_logger_get_default_io_unit()
     119          26 :       lattice = cell%hmat*angstrom
     120             : 
     121             :       ! change matrix to one d array
     122           8 :       DO i = 1, 3
     123          26 :          dpmd_cell((i - 1)*3 + 1:(i - 1)*3 + 3) = lattice(:, i)
     124             :       END DO
     125             : 
     126             :       ! get deepmd_data to save force, virial info
     127           2 :       CALL fist_nonbond_env_get(fist_nonbond_env, deepmd_data=deepmd_data)
     128           2 :       IF (.NOT. ASSOCIATED(deepmd_data)) THEN
     129          28 :          ALLOCATE (deepmd_data)
     130           2 :          CALL fist_nonbond_env_set(fist_nonbond_env, deepmd_data=deepmd_data)
     131           2 :          NULLIFY (deepmd_data%use_indices, deepmd_data%force)
     132           2 :          deepmd_data%model = deepmd_model_load(filename=pot%set(1)%deepmd%deepmd_file_name)
     133             :       END IF
     134           2 :       IF (ASSOCIATED(deepmd_data%force)) THEN
     135           0 :          IF (SIZE(deepmd_data%force, 2) /= n_atoms_use) THEN
     136           0 :             DEALLOCATE (deepmd_data%force, deepmd_data%use_indices)
     137             :          END IF
     138             :       END IF
     139           2 :       IF (.NOT. ASSOCIATED(deepmd_data%force)) THEN
     140           6 :          ALLOCATE (deepmd_data%force(3, n_atoms_use))
     141           4 :          ALLOCATE (deepmd_data%use_indices(n_atoms_use))
     142             :       END IF
     143             : 
     144             :       CALL deepmd_model_compute(model=deepmd_data%model, &
     145             :                                 natom=dpmd_natoms, &
     146             :                                 coord=dpmd_coord, &
     147             :                                 atype=dpmd_atype, &
     148             :                                 cell=dpmd_cell, &
     149             :                                 energy=pot_deepmd, &
     150             :                                 force=dpmd_force, &
     151             :                                 virial=dpmd_virial, &
     152             :                                 atomic_energy=dpmd_atomic_energy, &
     153           2 :                                 atomic_virial=dpmd_atomic_virial)
     154             : 
     155             :       ! convert units
     156           2 :       pot_deepmd = pot_deepmd/evolt
     157          50 :       dpmd_force = dpmd_force/(evolt/angstrom)
     158          20 :       dpmd_virial = dpmd_virial/evolt
     159             : 
     160             :       ! account for double counting from multiple MPI processes
     161           2 :       IF (PRESENT(para_env)) THEN
     162           2 :          pot_deepmd = pot_deepmd/REAL(para_env%num_pe, dp)
     163          50 :          dpmd_force = dpmd_force/REAL(para_env%num_pe, dp)
     164          20 :          dpmd_virial = dpmd_virial/REAL(para_env%num_pe, dp)
     165             :       END IF
     166             : 
     167             :       ! save force, virial info
     168             :       iat_use = 0
     169          18 :       DO iat = 1, n_atoms
     170          18 :          IF (use_atom(iat)) THEN
     171          16 :             iat_use = iat_use + 1
     172          16 :             deepmd_data%use_indices(iat_use) = iat
     173          64 :             deepmd_data%force(1:3, iat_use) = dpmd_force((iat_use - 1)*3 + 1:(iat_use - 1)*3 + 3)
     174             :          END IF
     175             :       END DO
     176           8 :       DO i = 1, 3
     177          26 :          deepmd_data%virial(1:3, i) = dpmd_virial((i - 1)*3 + 1:(i - 1)*3 + 3)
     178             :       END DO
     179           0 :       DEALLOCATE (use_atom, use_atom_type, dpmd_coord, dpmd_force, &
     180           0 :                   dpmd_virial, dpmd_atomic_energy, dpmd_atomic_virial, &
     181           2 :                   dpmd_cell, dpmd_atype)
     182             : 
     183           2 :       CALL timestop(handle)
     184           4 :    END SUBROUTINE deepmd_energy_store_force_virial
     185             : 
     186             : ! **************************************************************************************************
     187             : !> \brief ...
     188             : !> \param fist_nonbond_env ...
     189             : !> \param force ...
     190             : !> \param pv_nonbond ...
     191             : !> \param use_virial ...
     192             : ! **************************************************************************************************
     193           2 :    SUBROUTINE deepmd_add_force_virial(fist_nonbond_env, force, pv_nonbond, use_virial)
     194             :       TYPE(fist_nonbond_env_type), POINTER               :: fist_nonbond_env
     195             :       REAL(KIND=dp)                                      :: force(:, :), pv_nonbond(3, 3)
     196             :       LOGICAL, OPTIONAL                                  :: use_virial
     197             : 
     198             :       CHARACTER(LEN=*), PARAMETER :: routineN = 'deepmd_add_force_virial'
     199             : 
     200             :       INTEGER                                            :: handle, iat, iat_use
     201             :       TYPE(deepmd_data_type), POINTER                    :: deepmd_data
     202             : 
     203           2 :       CALL timeset(routineN, handle)
     204             : 
     205           2 :       CALL fist_nonbond_env_get(fist_nonbond_env, deepmd_data=deepmd_data)
     206             : 
     207           2 :       IF (.NOT. ASSOCIATED(deepmd_data)) RETURN
     208             : 
     209          18 :       DO iat_use = 1, SIZE(deepmd_data%use_indices)
     210          16 :          iat = deepmd_data%use_indices(iat_use)
     211          16 :          CPASSERT(iat >= 1 .AND. iat <= SIZE(force, 2))
     212          66 :          force(1:3, iat) = force(1:3, iat) + deepmd_data%force(1:3, iat_use)
     213             :       END DO
     214             : 
     215           2 :       IF (use_virial) THEN
     216           0 :          pv_nonbond = pv_nonbond + deepmd_data%virial
     217             :       END IF
     218             : 
     219           2 :       CALL timestop(handle)
     220             :    END SUBROUTINE deepmd_add_force_virial
     221             : 
     222             : END MODULE manybody_deepmd