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