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             : ! **************************************************************************************************
       9             : !> \brief Gaussian Process implementation
      10             : !> \author Ole Schuett
      11             : ! **************************************************************************************************
      12             : MODULE pao_ml_gaussprocess
      13             :    USE kinds,                           ONLY: dp
      14             :    USE pao_types,                       ONLY: pao_env_type,&
      15             :                                               training_matrix_type
      16             : #include "./base/base_uses.f90"
      17             : 
      18             :    IMPLICIT NONE
      19             : 
      20             :    PRIVATE
      21             : 
      22             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_ml_gaussprocess'
      23             : 
      24             :    PUBLIC ::pao_ml_gp_train, pao_ml_gp_predict, pao_ml_gp_gradient
      25             : 
      26             : CONTAINS
      27             : 
      28             : ! **************************************************************************************************
      29             : !> \brief Builds the covariance matrix
      30             : !> \param pao ...
      31             : ! **************************************************************************************************
      32          10 :    SUBROUTINE pao_ml_gp_train(pao)
      33             :       TYPE(pao_env_type), POINTER                        :: pao
      34             : 
      35             :       INTEGER                                            :: i, ikind, info, j, npoints
      36          10 :       REAL(dp), DIMENSION(:), POINTER                    :: idescr, jdescr
      37             :       TYPE(training_matrix_type), POINTER                :: training_matrix
      38             : 
      39             :       ! TODO this could be parallelized over ranks
      40          22 :       DO ikind = 1, SIZE(pao%ml_training_matrices)
      41          12 :          training_matrix => pao%ml_training_matrices(ikind)
      42          12 :          npoints = SIZE(training_matrix%inputs, 2) ! number of points
      43          12 :          CPASSERT(SIZE(training_matrix%outputs, 2) == npoints)
      44          12 :          IF (npoints == 0) CYCLE ! have no training data
      45             : 
      46          15 :          IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Building covariance matrix for kind: ", &
      47          10 :             TRIM(training_matrix%kindname), " from ", npoints, "training points."
      48             : 
      49             :          ! build covariance matrix
      50          40 :          ALLOCATE (training_matrix%GP(npoints, npoints))
      51          46 :          DO i = 1, npoints
      52         132 :          DO j = i, npoints
      53          86 :             idescr => training_matrix%inputs(:, i)
      54          86 :             jdescr => training_matrix%inputs(:, j)
      55          86 :             training_matrix%GP(i, j) = kernel(pao%gp_scale, idescr, jdescr)
      56         122 :             training_matrix%GP(j, i) = training_matrix%GP(i, j)
      57             :          END DO
      58             :          END DO
      59             : 
      60             :          ! add noise of training data
      61          46 :          DO i = 1, npoints
      62          46 :             training_matrix%GP(i, i) = training_matrix%GP(i, i) + pao%gp_noise_var**2
      63             :          END DO
      64             : 
      65             :          ! compute cholesky decomposition of covariance matrix
      66          10 :          CALL dpotrf("U", npoints, training_matrix%GP, npoints, info)
      67          20 :          CPASSERT(info == 0)
      68             :       END DO
      69             : 
      70          10 :    END SUBROUTINE pao_ml_gp_train
      71             : 
      72             : ! **************************************************************************************************
      73             : !> \brief Uses covariance matrix to make prediction
      74             : !> \param pao ...
      75             : !> \param ikind ...
      76             : !> \param descriptor ...
      77             : !> \param output ...
      78             : !> \param variance ...
      79             : ! **************************************************************************************************
      80          82 :    SUBROUTINE pao_ml_gp_predict(pao, ikind, descriptor, output, variance)
      81             :       TYPE(pao_env_type), POINTER                        :: pao
      82             :       INTEGER, INTENT(IN)                                :: ikind
      83             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: descriptor
      84             :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: output
      85             :       REAL(dp), INTENT(OUT)                              :: variance
      86             : 
      87             :       INTEGER                                            :: i, info, npoints
      88          82 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: cov, weights
      89             :       TYPE(training_matrix_type), POINTER                :: training_matrix
      90             : 
      91          82 :       training_matrix => pao%ml_training_matrices(ikind)
      92          82 :       npoints = SIZE(training_matrix%outputs, 2)
      93             : 
      94             :       ! calculate covariances between descriptor and training-points
      95         246 :       ALLOCATE (cov(npoints))
      96         406 :       DO i = 1, npoints
      97         406 :          cov(i) = kernel(pao%gp_scale, descriptor, training_matrix%inputs(:, i))
      98             :       END DO
      99             : 
     100             :       ! calculate weights
     101         164 :       ALLOCATE (weights(npoints))
     102         406 :       weights(:) = cov(:)
     103          82 :       CALL DPOTRS("U", npoints, 1, training_matrix%GP, npoints, weights, npoints, info)
     104          82 :       CPASSERT(info == 0)
     105             : 
     106             :       ! calculate predicted output
     107         656 :       output = 0.0_dp
     108         406 :       DO i = 1, npoints
     109        2674 :          output(:) = output + weights(i)*training_matrix%outputs(:, i)
     110             :       END DO
     111             : 
     112             :       ! calculate prediction's variance
     113         406 :       variance = kernel(pao%gp_scale, descriptor, descriptor) - DOT_PRODUCT(weights, cov)
     114             : 
     115          82 :       IF (variance < 0.0_dp) &
     116           0 :          CPABORT("PAO gaussian process found negative variance")
     117             : 
     118          82 :       DEALLOCATE (cov, weights)
     119          82 :    END SUBROUTINE pao_ml_gp_predict
     120             : 
     121             : ! **************************************************************************************************
     122             : !> \brief Calculate gradient of Gaussian process
     123             : !> \param pao ...
     124             : !> \param ikind ...
     125             : !> \param descriptor ...
     126             : !> \param outer_deriv ...
     127             : !> \param gradient ...
     128             : ! **************************************************************************************************
     129          10 :    SUBROUTINE pao_ml_gp_gradient(pao, ikind, descriptor, outer_deriv, gradient)
     130             :       TYPE(pao_env_type), POINTER                        :: pao
     131             :       INTEGER, INTENT(IN)                                :: ikind
     132             :       REAL(dp), DIMENSION(:), INTENT(IN), TARGET         :: descriptor
     133             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: outer_deriv
     134             :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: gradient
     135             : 
     136             :       INTEGER                                            :: i, info, npoints
     137          10 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: cov_deriv, weights_deriv
     138          20 :       REAL(dp), DIMENSION(SIZE(descriptor))              :: kg
     139             :       TYPE(training_matrix_type), POINTER                :: training_matrix
     140             : 
     141          10 :       training_matrix => pao%ml_training_matrices(ikind)
     142          10 :       npoints = SIZE(training_matrix%outputs, 2)
     143             : 
     144             :       ! calculate derivative of weights
     145          30 :       ALLOCATE (weights_deriv(npoints))
     146          46 :       DO i = 1, npoints
     147         298 :          weights_deriv(i) = SUM(outer_deriv*training_matrix%outputs(:, i))
     148             :       END DO
     149             : 
     150             :       ! calculate derivative of covariances
     151          20 :       ALLOCATE (cov_deriv(npoints))
     152          46 :       cov_deriv(:) = weights_deriv(:)
     153          10 :       CALL DPOTRS("U", npoints, 1, training_matrix%GP, npoints, cov_deriv, npoints, info)
     154          10 :       CPASSERT(info == 0)
     155             : 
     156             :       ! calculate derivative of kernel
     157         916 :       gradient(:) = 0.0_dp
     158          46 :       DO i = 1, npoints
     159          36 :          kg = kernel_grad(pao%gp_scale, descriptor, training_matrix%inputs(:, i))
     160        3082 :          gradient(:) = gradient(:) + kg(:)*cov_deriv(i)
     161             :       END DO
     162             : 
     163          10 :       DEALLOCATE (cov_deriv, weights_deriv)
     164          10 :    END SUBROUTINE pao_ml_gp_gradient
     165             : 
     166             : ! **************************************************************************************************
     167             : !> \brief Gaussian kernel used to measure covariance between two descriptors.
     168             : !> \param scale ...
     169             : !> \param descr1 ...
     170             : !> \param descr2 ...
     171             : !> \return ...
     172             : ! **************************************************************************************************
     173         492 :    PURE FUNCTION kernel(scale, descr1, descr2) RESULT(cov)
     174             :       REAL(dp), INTENT(IN)                               :: scale
     175             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: descr1, descr2
     176             :       REAL(dp)                                           :: cov
     177             : 
     178             :       REAL(dp)                                           :: fdist2
     179         492 :       REAL(dp), DIMENSION(SIZE(descr1))                  :: diff
     180             : 
     181       48036 :       diff = descr1 - descr2
     182       48036 :       fdist2 = SUM((diff/scale)**2)
     183         492 :       cov = EXP(-fdist2/2.0_dp)
     184         492 :    END FUNCTION kernel
     185             : 
     186             : ! **************************************************************************************************
     187             : !> \brief Gradient of Gaussian kernel wrt descr1
     188             : !> \param scale ...
     189             : !> \param descr1 ...
     190             : !> \param descr2 ...
     191             : !> \return ...
     192             : ! **************************************************************************************************
     193          72 :    PURE FUNCTION kernel_grad(scale, descr1, descr2) RESULT(grad)
     194             :       REAL(dp), INTENT(IN)                               :: scale
     195             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: descr1, descr2
     196             :       REAL(dp), DIMENSION(SIZE(descr1))                  :: grad
     197             : 
     198             :       REAL(dp)                                           :: cov, fdist2
     199          36 :       REAL(dp), DIMENSION(SIZE(descr1))                  :: diff
     200             : 
     201        3072 :       diff = descr1 - descr2
     202        3072 :       fdist2 = SUM((diff/scale)**2)
     203          36 :       cov = EXP(-fdist2/2.0_dp)
     204        3072 :       grad(:) = cov*(-diff/scale**2)
     205             : 
     206          36 :    END FUNCTION kernel_grad
     207             : 
     208             : END MODULE pao_ml_gaussprocess