pao_ml_descriptor.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief Feature vectors for describing chemical environments in a rotationally invariant fashion.
!> \author Ole Schuett
! **************************************************************************************************
MODULE pao_ml_descriptor
   USE atomic_kind_types,               ONLY: get_atomic_kind
   USE basis_set_types,                 ONLY: gto_basis_set_type
   USE cell_types,                      ONLY: cell_type,&
                                              pbc
   USE kinds,                           ONLY: dp
   USE mathconstants,                   ONLY: fourpi,&
                                              rootpi
   USE mathlib,                         ONLY: diamat_all
   USE pao_input,                       ONLY: pao_ml_desc_overlap,&
                                              pao_ml_desc_pot,&
                                              pao_ml_desc_r12
   USE pao_potentials,                  ONLY: pao_calc_gaussian
   USE pao_types,                       ONLY: pao_env_type
   USE particle_types,                  ONLY: particle_type
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              pao_descriptor_type,&
                                              qs_kind_type
   USE util,                            ONLY: sort
#include "./base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_ml_descriptor'
 
   PUBLIC :: pao_ml_calc_descriptor
 
CONTAINS
 
! **************************************************************************************************
!> \brief Calculates a descriptor for chemical environment of given atom
!> \param pao ...
!> \param particle_set ...
!> \param qs_kind_set ...
!> \param cell ...
!> \param iatom ...
!> \param descriptor ...
!> \param descr_grad ...
!> \param forces ...
! **************************************************************************************************
   SUBROUTINE pao_ml_calc_descriptor(pao, particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      TYPE(pao_env_type), POINTER                        :: pao
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(cell_type), POINTER                           :: cell
      INTEGER, INTENT(IN)                                :: iatom
      REAL(dp), ALLOCATABLE, DIMENSION(:), OPTIONAL      :: descriptor
      REAL(dp), DIMENSION(:), INTENT(IN), OPTIONAL       :: descr_grad
      REAL(dp), DIMENSION(:, :), INTENT(INOUT), OPTIONAL :: forces
 
      CHARACTER(len=*), PARAMETER :: routinen = 'pao_ml_calc_descriptor'
 
      INTEGER                                            :: handle
 
      CALL timeset(routinen, handle)
 
      cpassert(PRESENT(forces) .EQV. PRESENT(descr_grad))
 
      SELECT CASE (pao%ml_descriptor)
      CASE (pao_ml_desc_pot)
         CALL calc_descriptor_pot(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      CASE (pao_ml_desc_overlap)
         CALL calc_descriptor_overlap(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      CASE (pao_ml_desc_r12)
         CALL calc_descriptor_r12(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      CASE DEFAULT
         cpabort("PAO: unknown descriptor")
      END SELECT
 
      CALL timestop(handle)
   END SUBROUTINE pao_ml_calc_descriptor
 
! **************************************************************************************************
!> \brief Calculates a descriptor based on the eigenvalues of V_neighbors
!> \param particle_set ...
!> \param qs_kind_set ...
!> \param cell ...
!> \param iatom ...
!> \param descriptor ...
!> \param descr_grad ...
!> \param forces ...
! **************************************************************************************************
   SUBROUTINE calc_descriptor_pot(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(cell_type), POINTER                           :: cell
      INTEGER, INTENT(IN)                                :: iatom
      REAL(dp), ALLOCATABLE, DIMENSION(:), OPTIONAL      :: descriptor
      REAL(dp), DIMENSION(:), INTENT(IN), OPTIONAL       :: descr_grad
      REAL(dp), DIMENSION(:, :), INTENT(INOUT), OPTIONAL :: forces
 
      CHARACTER(len=*), PARAMETER :: routinen = 'calc_descriptor_pot'
 
      INTEGER                                            :: handle, i, idesc, ikind, jatom, jkind, &
                                                            k, n, natoms, ndesc
      REAL(dp)                                           :: beta, w, weight
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: v_evals
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: block_m, block_v, v_evecs
      REAL(dp), ALLOCATABLE, DIMENSION(:, :, :)          :: block_d
      REAL(dp), DIMENSION(3)                             :: ra, rab, rb
      TYPE(gto_basis_set_type), POINTER                  :: basis_set
      TYPE(pao_descriptor_type), DIMENSION(:), POINTER   :: pao_descriptors
 
      CALL timeset(routinen, handle)
 
      CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
      CALL get_qs_kind(qs_kind_set(ikind), basis_set=basis_set, pao_descriptors=pao_descriptors)
      n = basis_set%nsgf
      natoms = SIZE(particle_set)
      ndesc = SIZE(pao_descriptors)
      IF (ndesc == 0) cpabort("No PAO_DESCRIPTOR section found")
 
      ALLOCATE (block_v(n, n), v_evecs(n, n), v_evals(n))
      IF (PRESENT(descriptor)) ALLOCATE (descriptor(n*ndesc))
      IF (PRESENT(forces)) ALLOCATE (block_d(n, n, 3), block_m(n, n))
 
      DO idesc = 1, ndesc
 
         ! construct matrix V_block from neighboring atoms
         block_v = 0.0_dp
         DO jatom = 1, natoms
            IF (jatom == iatom) cycle
            ra = particle_set(iatom)%r
            rb = particle_set(jatom)%r
            rab = pbc(ra, rb, cell)
            CALL get_atomic_kind(particle_set(jatom)%atomic_kind, kind_number=jkind)
            CALL get_qs_kind(qs_kind_set(jkind), pao_descriptors=pao_descriptors)
            IF (SIZE(pao_descriptors) /= ndesc) THEN
               cpabort("Not all KINDs have the same number of PAO_DESCRIPTOR sections")
            END IF
            weight = pao_descriptors(idesc)%weight
            beta = pao_descriptors(idesc)%beta
            CALL pao_calc_gaussian(basis_set, block_v=block_v, rab=rab, lpot=0, beta=beta, weight=weight)
         END DO
 
         ! diagonalize block_V
         v_evecs(:, :) = block_v(:, :)
         CALL diamat_all(v_evecs, v_evals)
 
         ! use eigenvalues of V_block as descriptor
         IF (PRESENT(descriptor)) THEN
            descriptor((idesc - 1)*n + 1:idesc*n) = v_evals(:)
         END IF
 
         ! FORCES ----------------------------------------------------------------------------------
         IF (PRESENT(forces)) THEN
            cpassert(PRESENT(descr_grad))
            block_m = 0.0_dp
            DO k = 1, n
               w = descr_grad((idesc - 1)*n + k)
               block_m(:, :) = block_m(:, :) + w*matmul(v_evecs(:, k:k), transpose(v_evecs(:, k:k)))
            END DO
            DO jatom = 1, natoms
               IF (jatom == iatom) cycle
               ra = particle_set(iatom)%r
               rb = particle_set(jatom)%r
               rab = pbc(ra, rb, cell)
               CALL get_atomic_kind(particle_set(jatom)%atomic_kind, kind_number=jkind)
               CALL get_qs_kind(qs_kind_set(jkind), pao_descriptors=pao_descriptors)
               weight = pao_descriptors(idesc)%weight
               beta = pao_descriptors(idesc)%beta
               block_d = 0.0_dp
               CALL pao_calc_gaussian(basis_set, block_d=block_d, rab=rab, lpot=0, beta=beta, weight=weight)
               DO i = 1, 3
                  forces(iatom, i) = forces(iatom, i) - sum(block_m*block_d(:, :, i))
                  forces(jatom, i) = forces(jatom, i) + sum(block_m*block_d(:, :, i))
               END DO
            END DO
         END IF
 
      END DO
 
      CALL timestop(handle)
   END SUBROUTINE calc_descriptor_pot
 
! **************************************************************************************************
!> \brief Calculates a descriptor based on the eigenvalues of local overlap matrix
!> \param particle_set ...
!> \param qs_kind_set ...
!> \param cell ...
!> \param iatom ...
!> \param descriptor ...
!> \param descr_grad ...
!> \param forces ...
! **************************************************************************************************
   SUBROUTINE calc_descriptor_overlap(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(cell_type), POINTER                           :: cell
      INTEGER, INTENT(IN)                                :: iatom
      REAL(dp), ALLOCATABLE, DIMENSION(:), OPTIONAL      :: descriptor
      REAL(dp), DIMENSION(:), INTENT(IN), OPTIONAL       :: descr_grad
      REAL(dp), DIMENSION(:, :), INTENT(INOUT), OPTIONAL :: forces
 
      CHARACTER(len=*), PARAMETER :: routinen = 'calc_descriptor_overlap'
 
      INTEGER                                            :: handle, idesc, ikind, j, jatom, jkind, &
                                                            k, katom, kkind, n, natoms, ndesc
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: neighbor_order
      REAL(dp) :: beta_sum, deriv, exponent, integral, jbeta, jweight, kbeta, kweight, &
         normalization, rij2, rik2, rjk2, sbeta, screening_radius, screening_volume, w
      REAL(dp), ALLOCATABLE, DIMENSION(:)                :: s_evals
      REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: block_m, block_s, s_evecs
      REAL(dp), DIMENSION(3)                             :: ri, rij, rik, rj, rjk, rk
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:)           :: neighbor_dist
      TYPE(pao_descriptor_type), DIMENSION(:), POINTER   :: ipao_descriptors, jpao_descriptors, &
                                                            kpao_descriptors
 
      CALL timeset(routinen, handle)
 
      CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
      CALL get_qs_kind(qs_kind_set(ikind), pao_descriptors=ipao_descriptors)
 
      natoms = SIZE(particle_set)
      ndesc = SIZE(ipao_descriptors)
      IF (ndesc == 0) cpabort("No PAO_DESCRIPTOR section found")
 
      ! determine largest screening radius
      screening_radius = 0.0_dp
      DO idesc = 1, ndesc
         screening_radius = max(screening_radius, ipao_descriptors(idesc)%screening_radius)
      END DO
 
      ! estimate maximum number of neighbors within screening
      screening_volume = fourpi/3.0_dp*screening_radius**3
      n = int(screening_volume/35.0_dp) ! rule of thumb
 
      ALLOCATE (block_s(n, n), s_evals(n), s_evecs(n, n))
      IF (PRESENT(descriptor)) ALLOCATE (descriptor(n*ndesc))
      IF (PRESENT(forces)) ALLOCATE (block_m(n, n))
 
      !find neighbors
      !TODO: this is a quadratic algorithm, use a neighbor-list instead
      ALLOCATE (neighbor_dist(natoms), neighbor_order(natoms))
      ri = particle_set(iatom)%r
      DO jatom = 1, natoms
         rj = particle_set(jatom)%r
         rij = pbc(ri, rj, cell)
         neighbor_dist(jatom) = sqrt(sum(rij**2))
      END DO
      CALL sort(neighbor_dist, natoms, neighbor_order)
      cpassert(neighbor_order(1) == iatom) !central atom should be closesd to itself
 
      ! check if N was chosen large enough
      IF (natoms > n) THEN
         IF (neighbor_dist(n + 1) < screening_radius) THEN
            cpabort("PAO heuristic for descriptor size broke down")
         END IF
      END IF
 
      DO idesc = 1, ndesc
         sbeta = ipao_descriptors(idesc)%screening
 
         ! construct matrix S_block from neighboring atoms
         block_s = 0.0_dp
         DO j = 1, min(natoms, n)
         DO k = 1, min(natoms, n)
            jatom = neighbor_order(j)
            katom = neighbor_order(k)
 
            ! get weigths and betas
            CALL get_atomic_kind(particle_set(jatom)%atomic_kind, kind_number=jkind)
            CALL get_qs_kind(qs_kind_set(jkind), pao_descriptors=jpao_descriptors)
            CALL get_atomic_kind(particle_set(katom)%atomic_kind, kind_number=kkind)
            CALL get_qs_kind(qs_kind_set(kkind), pao_descriptors=kpao_descriptors)
            IF (SIZE(jpao_descriptors) /= ndesc .OR. SIZE(kpao_descriptors) /= ndesc) THEN
               cpabort("Not all KINDs have the same number of PAO_DESCRIPTOR sections")
            END IF
            jweight = jpao_descriptors(idesc)%weight
            jbeta = jpao_descriptors(idesc)%beta
            kweight = kpao_descriptors(idesc)%weight
            kbeta = kpao_descriptors(idesc)%beta
            beta_sum = sbeta + jbeta + kbeta
 
            ! get distances
            rj = particle_set(jatom)%r
            rk = particle_set(katom)%r
            rij = pbc(ri, rj, cell)
            rik = pbc(ri, rk, cell)
            rjk = pbc(rj, rk, cell)
            rij2 = sum(rij**2)
            rik2 = sum(rik**2)
            rjk2 = sum(rjk**2)
 
            ! calculate integral over three Gaussians
            exponent = -(sbeta*jbeta*rij2 + sbeta*kbeta*rik2 + jbeta*kbeta*rjk2)/beta_sum
            integral = exp(exponent)*rootpi/sqrt(beta_sum)
            normalization = sqrt(jbeta*kbeta)/rootpi**2
            block_s(j, k) = jweight*kweight*normalization*integral
         END DO
         END DO
 
         ! diagonalize V_block
         s_evecs(:, :) = block_s(:, :)
         CALL diamat_all(s_evecs, s_evals)
 
         ! use eigenvalues of S_block as descriptor
         IF (PRESENT(descriptor)) THEN
            descriptor((idesc - 1)*n + 1:idesc*n) = s_evals(:)
         END IF
 
         ! FORCES ----------------------------------------------------------------------------------
         IF (PRESENT(forces)) THEN
            cpassert(PRESENT(descr_grad))
            block_m = 0.0_dp
            DO k = 1, n
               w = descr_grad((idesc - 1)*n + k)
               block_m(:, :) = block_m(:, :) + w*matmul(s_evecs(:, k:k), transpose(s_evecs(:, k:k)))
            END DO
 
            DO j = 1, min(natoms, n)
            DO k = 1, min(natoms, n)
               jatom = neighbor_order(j)
               katom = neighbor_order(k)
 
               ! get weigths and betas
               CALL get_atomic_kind(particle_set(jatom)%atomic_kind, kind_number=jkind)
               CALL get_qs_kind(qs_kind_set(jkind), pao_descriptors=jpao_descriptors)
               CALL get_atomic_kind(particle_set(katom)%atomic_kind, kind_number=kkind)
               CALL get_qs_kind(qs_kind_set(kkind), pao_descriptors=kpao_descriptors)
               jweight = jpao_descriptors(idesc)%weight
               jbeta = jpao_descriptors(idesc)%beta
               kweight = kpao_descriptors(idesc)%weight
               kbeta = kpao_descriptors(idesc)%beta
               beta_sum = sbeta + jbeta + kbeta
 
               ! get distances
               rj = particle_set(jatom)%r
               rk = particle_set(katom)%r
               rij = pbc(ri, rj, cell)
               rik = pbc(ri, rk, cell)
               rjk = pbc(rj, rk, cell)
               rij2 = sum(rij**2)
               rik2 = sum(rik**2)
               rjk2 = sum(rjk**2)
 
               ! calculate integral over three Gaussians
               exponent = -(sbeta*jbeta*rij2 + sbeta*kbeta*rik2 + jbeta*kbeta*rjk2)/beta_sum
               integral = exp(exponent)*rootpi/sqrt(beta_sum)
               normalization = sqrt(jbeta*kbeta)/rootpi**2
               deriv = 2.0_dp/beta_sum*block_m(j, k)
               w = jweight*kweight*normalization*integral*deriv
               forces(iatom, :) = forces(iatom, :) - sbeta*jbeta*rij*w
               forces(jatom, :) = forces(jatom, :) + sbeta*jbeta*rij*w
               forces(iatom, :) = forces(iatom, :) - sbeta*kbeta*rik*w
               forces(katom, :) = forces(katom, :) + sbeta*kbeta*rik*w
               forces(jatom, :) = forces(jatom, :) - jbeta*kbeta*rjk*w
               forces(katom, :) = forces(katom, :) + jbeta*kbeta*rjk*w
            END DO
            END DO
         END IF
      END DO
 
      CALL timestop(handle)
   END SUBROUTINE calc_descriptor_overlap
 
! **************************************************************************************************
!> \brief Calculates a descriptor based on distance between two atoms
!> \param particle_set ...
!> \param qs_kind_set ...
!> \param cell ...
!> \param iatom ...
!> \param descriptor ...
!> \param descr_grad ...
!> \param forces ...
! **************************************************************************************************
   SUBROUTINE calc_descriptor_r12(particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      TYPE(cell_type), POINTER                           :: cell
      INTEGER, INTENT(IN)                                :: iatom
      REAL(dp), ALLOCATABLE, DIMENSION(:), OPTIONAL      :: descriptor
      REAL(dp), DIMENSION(:), INTENT(IN), OPTIONAL       :: descr_grad
      REAL(dp), DIMENSION(:, :), INTENT(INOUT), OPTIONAL :: forces
 
      REAL(dp), DIMENSION(3)                             :: g, r1, r12, r2
 
      cpassert(SIZE(particle_set) == 2)
 
      mark_used(qs_kind_set)
      mark_used(iatom)
      mark_used(cell)
 
      r1 = particle_set(1)%r
      r2 = particle_set(2)%r
      r12 = pbc(r1, r2, cell)
 
      IF (PRESENT(descriptor)) THEN
         ALLOCATE (descriptor(1))
         descriptor(1) = sqrt(sum(r12**2))
      END IF
 
      IF (PRESENT(forces)) THEN
         cpassert(PRESENT(descr_grad))
         g = r12/sqrt(sum(r12**2))*descr_grad(1)
         forces(1, :) = forces(1, :) + g
         forces(2, :) = forces(2, :) - g
      END IF
   END SUBROUTINE calc_descriptor_r12
 
END MODULE pao_ml_descriptor