pao_param_equi.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief Equivariant parametrization
!> \author Ole Schuett
! **************************************************************************************************
MODULE pao_param_equi
   USE basis_set_types,                 ONLY: gto_basis_set_type
   USE cp_dbcsr_api,                    ONLY: &
        dbcsr_complete_redistribute, dbcsr_create, dbcsr_distribution_type, dbcsr_get_block_p, &
        dbcsr_get_info, dbcsr_iterator_blocks_left, dbcsr_iterator_next_block, &
        dbcsr_iterator_start, dbcsr_iterator_stop, dbcsr_iterator_type, dbcsr_p_type, &
        dbcsr_release, dbcsr_type
   USE cp_dbcsr_contrib,                ONLY: dbcsr_reserve_diag_blocks
   USE dm_ls_scf_types,                 ONLY: ls_mstruct_type,&
                                              ls_scf_env_type
   USE kinds,                           ONLY: dp
   USE mathlib,                         ONLY: diamat_all
   USE message_passing,                 ONLY: mp_comm_type
   USE pao_param_methods,               ONLY: pao_calc_grad_lnv_wrt_ab
   USE pao_potentials,                  ONLY: pao_guess_initial_potential
   USE pao_types,                       ONLY: pao_env_type
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              qs_kind_type
#include "./base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_param_equi'
 
   PUBLIC :: pao_param_init_equi, pao_param_finalize_equi, pao_calc_ab_equi
   PUBLIC :: pao_param_count_equi, pao_param_initguess_equi
 
CONTAINS
 
! **************************************************************************************************
!> \brief Initialize equivariant parametrization
!> \param pao ...
! **************************************************************************************************
   SUBROUTINE pao_param_init_equi(pao)
      TYPE(pao_env_type), POINTER                        :: pao
 
      IF (pao%precondition) &
         cpabort("PAO preconditioning not supported for selected parametrization.")
 
   END SUBROUTINE pao_param_init_equi
 
! **************************************************************************************************
!> \brief Finalize equivariant parametrization
! **************************************************************************************************
   SUBROUTINE pao_param_finalize_equi()
 
      ! Nothing to do.
 
   END SUBROUTINE pao_param_finalize_equi
 
! **************************************************************************************************
!> \brief Returns the number of parameters for given atomic kind
!> \param qs_env ...
!> \param ikind ...
!> \param nparams ...
! **************************************************************************************************
   SUBROUTINE pao_param_count_equi(qs_env, ikind, nparams)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      INTEGER, INTENT(IN)                                :: ikind
      INTEGER, INTENT(OUT)                               :: nparams
 
      INTEGER                                            :: pao_basis_size, pri_basis_size
      TYPE(gto_basis_set_type), POINTER                  :: basis_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
 
      CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set)
      CALL get_qs_kind(qs_kind_set(ikind), &
                       basis_set=basis_set, &
                       pao_basis_size=pao_basis_size)
      pri_basis_size = basis_set%nsgf
 
      nparams = pao_basis_size*pri_basis_size
 
   END SUBROUTINE pao_param_count_equi
 
! **************************************************************************************************
!> \brief Fills matrix_X with an initial guess
!> \param pao ...
!> \param qs_env ...
! **************************************************************************************************
   SUBROUTINE pao_param_initguess_equi(pao, qs_env)
      TYPE(pao_env_type), POINTER                        :: pao
      TYPE(qs_environment_type), POINTER                 :: qs_env
 
      CHARACTER(len=*), PARAMETER :: routinen = 'pao_param_initguess_equi'
 
      INTEGER                                            :: acol, arow, handle, i, iatom, m, n
      INTEGER, DIMENSION(:), POINTER                     :: blk_sizes_pao, blk_sizes_pri
      LOGICAL                                            :: found
      REAL(dp), DIMENSION(:), POINTER                    :: h_evals
      REAL(dp), DIMENSION(:, :), POINTER                 :: a, block_h0, block_n, block_n_inv, &
                                                            block_x, h, h_evecs, v0
      TYPE(dbcsr_iterator_type)                          :: iter
 
      CALL timeset(routinen, handle)
 
      CALL dbcsr_get_info(pao%matrix_Y, row_blk_size=blk_sizes_pri, col_blk_size=blk_sizes_pao)
 
!$OMP PARALLEL DEFAULT(NONE) SHARED(pao,qs_env,blk_sizes_pri,blk_sizes_pao) &
!$OMP PRIVATE(iter,arow,acol,iatom,n,m,i,found) &
!$OMP PRIVATE(block_X,block_H0,block_N,block_N_inv,A,H,H_evecs,H_evals,V0)
      CALL dbcsr_iterator_start(iter, pao%matrix_X)
      DO WHILE (dbcsr_iterator_blocks_left(iter))
         CALL dbcsr_iterator_next_block(iter, arow, acol, block_x)
         iatom = arow; cpassert(arow == acol)
 
         CALL dbcsr_get_block_p(matrix=pao%matrix_H0, row=iatom, col=iatom, block=block_h0, found=found)
         CALL dbcsr_get_block_p(matrix=pao%matrix_N_diag, row=iatom, col=iatom, block=block_n, found=found)
         CALL dbcsr_get_block_p(matrix=pao%matrix_N_inv_diag, row=iatom, col=iatom, block=block_n_inv, found=found)
         cpassert(ASSOCIATED(block_h0) .AND. ASSOCIATED(block_n) .AND. ASSOCIATED(block_n_inv))
 
         n = blk_sizes_pri(iatom) ! size of primary basis
         m = blk_sizes_pao(iatom) ! size of pao basis
 
         ALLOCATE (v0(n, n))
         CALL pao_guess_initial_potential(qs_env, iatom, v0)
 
         ! construct H
         ALLOCATE (h(n, n))
         h = matmul(matmul(block_n, block_h0 + v0), block_n) ! transform into orthonormal basis
 
         ! diagonalize H
         ALLOCATE (h_evecs(n, n), h_evals(n))
         h_evecs = h
         CALL diamat_all(h_evecs, h_evals)
 
         ! use first m eigenvectors as initial guess
         ALLOCATE (a(n, m))
         a = matmul(block_n_inv, h_evecs(:, 1:m))
 
         ! normalize vectors
         DO i = 1, m
            a(:, i) = a(:, i)/norm2(a(:, i))
         END DO
 
         block_x = reshape(a, (/n*m, 1/))
         DEALLOCATE (h, v0, a, h_evecs, h_evals)
 
      END DO
      CALL dbcsr_iterator_stop(iter)
!$OMP END PARALLEL
 
      CALL timestop(handle)
 
   END SUBROUTINE pao_param_initguess_equi
 
! **************************************************************************************************
!> \brief Takes current matrix_X and calculates the matrices A and B.
!> \param pao ...
!> \param qs_env ...
!> \param ls_scf_env ...
!> \param gradient ...
!> \param penalty ...
! **************************************************************************************************
   SUBROUTINE pao_calc_ab_equi(pao, qs_env, ls_scf_env, gradient, penalty)
      TYPE(pao_env_type), POINTER                        :: pao
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(ls_scf_env_type), TARGET                      :: ls_scf_env
      LOGICAL, INTENT(IN)                                :: gradient
      REAL(dp), INTENT(INOUT), OPTIONAL                  :: penalty
 
      CHARACTER(len=*), PARAMETER                        :: routinen = 'pao_calc_AB_equi'
 
      INTEGER                                            :: acol, arow, handle, i, iatom, j, k, m, n
      LOGICAL                                            :: found
      REAL(dp)                                           :: denom, w
      REAL(dp), DIMENSION(:), POINTER                    :: anna_evals
      REAL(dp), DIMENSION(:, :), POINTER                 :: anna, anna_evecs, anna_inv, block_a, &
                                                            block_b, block_g, block_ma, block_mb, &
                                                            block_n, block_x, d, g, m1, m2, m3, &
                                                            m4, m5, nn
      TYPE(dbcsr_distribution_type)                      :: main_dist
      TYPE(dbcsr_iterator_type)                          :: iter
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_s
      TYPE(dbcsr_type)                                   :: matrix_g_nondiag, matrix_ma, matrix_mb, &
                                                            matrix_x_nondiag
      TYPE(ls_mstruct_type), POINTER                     :: ls_mstruct
      TYPE(mp_comm_type)                                 :: group
 
      CALL timeset(routinen, handle)
      ls_mstruct => ls_scf_env%ls_mstruct
 
      IF (gradient) THEN
         CALL pao_calc_grad_lnv_wrt_ab(qs_env, ls_scf_env, matrix_ma, matrix_mb)
      END IF
 
      ! Redistribute matrix_X from diag_distribution to distribution of matrix_s.
      CALL get_qs_env(qs_env, matrix_s=matrix_s)
      CALL dbcsr_get_info(matrix=matrix_s(1)%matrix, distribution=main_dist)
      CALL dbcsr_create(matrix_x_nondiag, &
                        name="PAO matrix_X_nondiag", &
                        dist=main_dist, &
                        template=pao%matrix_X)
      CALL dbcsr_reserve_diag_blocks(matrix_x_nondiag)
      CALL dbcsr_complete_redistribute(pao%matrix_X, matrix_x_nondiag)
 
      ! Compuation of matrix_G uses distr. of matrix_s, afterwards we redistribute to diag_distribution.
      IF (gradient) THEN
         CALL dbcsr_create(matrix_g_nondiag, &
                           name="PAO matrix_G_nondiag", &
                           dist=main_dist, &
                           template=pao%matrix_G)
         CALL dbcsr_reserve_diag_blocks(matrix_g_nondiag)
      END IF
 
!$OMP PARALLEL DEFAULT(NONE) &
!$OMP SHARED(pao,ls_mstruct,matrix_X_nondiag,matrix_G_nondiag,matrix_Ma,matrix_Mb,gradient,penalty) &
!$OMP PRIVATE(iter,arow,acol,iatom,found,n,m,w,i,j,k,denom) &
!$OMP PRIVATE(NN,ANNA,ANNA_evals,ANNA_evecs,ANNA_inv,D,G,M1,M2,M3,M4,M5) &
!$OMP PRIVATE(block_X,block_A,block_B,block_N,block_Ma, block_Mb, block_G)
      CALL dbcsr_iterator_start(iter, matrix_x_nondiag)
      DO WHILE (dbcsr_iterator_blocks_left(iter))
         CALL dbcsr_iterator_next_block(iter, arow, acol, block_x)
         iatom = arow; cpassert(arow == acol)
         CALL dbcsr_get_block_p(matrix=ls_mstruct%matrix_A, row=iatom, col=iatom, block=block_a, found=found)
         cpassert(ASSOCIATED(block_a))
         CALL dbcsr_get_block_p(matrix=ls_mstruct%matrix_B, row=iatom, col=iatom, block=block_b, found=found)
         cpassert(ASSOCIATED(block_b))
         CALL dbcsr_get_block_p(matrix=pao%matrix_N, row=iatom, col=iatom, block=block_n, found=found)
         cpassert(ASSOCIATED(block_n))
 
         n = SIZE(block_a, 1) ! size of primary basis
         m = SIZE(block_a, 2) ! size of pao basis
         block_a = reshape(block_x, (/n, m/))
 
         ! restrain pao basis vectors to unit norm
         IF (PRESENT(penalty)) THEN
            DO i = 1, m
               w = 1.0_dp - sum(block_a(:, i)**2)
               penalty = penalty + pao%penalty_strength*w**2
            END DO
         END IF
 
         ALLOCATE (nn(n, n), anna(m, m))
         nn = matmul(block_n, block_n) ! it's actually S^{-1}
         anna = matmul(matmul(transpose(block_a), nn), block_a)
 
         ! diagonalize ANNA
         ALLOCATE (anna_evecs(m, m), anna_evals(m))
         anna_evecs(:, :) = anna
         CALL diamat_all(anna_evecs, anna_evals)
         IF (minval(abs(anna_evals)) < 1e-10_dp) cpabort("PAO basis singualar.")
 
         ! build ANNA_inv
         ALLOCATE (anna_inv(m, m))
         anna_inv(:, :) = 0.0_dp
         DO k = 1, m
            w = 1.0_dp/anna_evals(k)
            DO i = 1, m
            DO j = 1, m
               anna_inv(i, j) = anna_inv(i, j) + w*anna_evecs(i, k)*anna_evecs(j, k)
            END DO
            END DO
         END DO
 
         !B = 1/S * A * 1/(A^T 1/S A)
         block_b = matmul(matmul(nn, block_a), anna_inv)
 
         ! TURNING POINT (if calc grad) ------------------------------------------
         IF (gradient) THEN
            CALL dbcsr_get_block_p(matrix=matrix_g_nondiag, row=iatom, col=iatom, block=block_g, found=found)
            cpassert(ASSOCIATED(block_g))
            CALL dbcsr_get_block_p(matrix=matrix_ma, row=iatom, col=iatom, block=block_ma, found=found)
            CALL dbcsr_get_block_p(matrix=matrix_mb, row=iatom, col=iatom, block=block_mb, found=found)
            ! don't check ASSOCIATED(block_M), it might have been filtered out.
 
            ALLOCATE (g(n, m))
            g(:, :) = 0.0_dp
 
            IF (PRESENT(penalty)) THEN
               DO i = 1, m
                  w = 1.0_dp - sum(block_a(:, i)**2)
                  g(:, i) = -4.0_dp*pao%penalty_strength*w*block_a(:, i)
               END DO
            END IF
 
            IF (ASSOCIATED(block_ma)) THEN
               g = g + block_ma
            END IF
 
            IF (ASSOCIATED(block_mb)) THEN
               g = g + matmul(matmul(nn, block_mb), anna_inv)
 
               ! calculate derivatives dAA_inv/ dAA
               ALLOCATE (d(m, m), m1(m, m), m2(m, m), m3(m, m), m4(m, m), m5(m, m))
 
               DO i = 1, m
               DO j = 1, m
                  denom = anna_evals(i) - anna_evals(j)
                  IF (i == j) THEN
                     d(i, i) = -1.0_dp/anna_evals(i)**2 ! diagonal elements
                  ELSE IF (abs(denom) > 1e-10_dp) THEN
                     d(i, j) = (1.0_dp/anna_evals(i) - 1.0_dp/anna_evals(j))/denom
                  ELSE
                     d(i, j) = -1.0_dp ! limit according to L'Hospital's rule
                  END IF
               END DO
               END DO
 
               m1 = matmul(matmul(transpose(block_a), nn), block_mb)
               m2 = matmul(matmul(transpose(anna_evecs), m1), anna_evecs)
               m3 = m2*d ! Hadamard product
               m4 = matmul(matmul(anna_evecs, m3), transpose(anna_evecs))
               m5 = 0.5_dp*(m4 + transpose(m4))
               g = g + 2.0_dp*matmul(matmul(nn, block_a), m5)
 
               DEALLOCATE (d, m1, m2, m3, m4, m5)
            END IF
 
            block_g = reshape(g, (/n*m, 1/))
            DEALLOCATE (g)
         END IF
 
         DEALLOCATE (nn, anna, anna_evecs, anna_evals, anna_inv)
      END DO
      CALL dbcsr_iterator_stop(iter)
!$OMP END PARALLEL
 
      ! sum penalty energies across ranks
      IF (PRESENT(penalty)) THEN
         CALL dbcsr_get_info(pao%matrix_X, group=group)
         CALL group%sum(penalty)
      END IF
 
      CALL dbcsr_release(matrix_x_nondiag)
 
      IF (gradient) THEN
         CALL dbcsr_complete_redistribute(matrix_g_nondiag, pao%matrix_G)
         CALL dbcsr_release(matrix_g_nondiag)
         CALL dbcsr_release(matrix_ma)
         CALL dbcsr_release(matrix_mb)
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE pao_calc_ab_equi
 
END MODULE pao_param_equi