LCOV - code coverage report
Current view: top level - src - pao_ml_neuralnet.F (source / functions) Coverage Total Hit
Test: CP2K Regtests (git:42dac4a) Lines: 99.1 % 108 107
Test Date: 2025-07-25 12:55:17 Functions: 100.0 % 5 5

            Line data    Source code
       1              : !--------------------------------------------------------------------------------------------------!
       2              : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3              : !   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
       4              : !                                                                                                  !
       5              : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6              : !--------------------------------------------------------------------------------------------------!
       7              : 
       8              : ! **************************************************************************************************
       9              : !> \brief Neural Network implementation
      10              : !> \author Ole Schuett
      11              : ! **************************************************************************************************
      12              : MODULE pao_ml_neuralnet
      13              :    USE kinds,                           ONLY: dp
      14              :    USE pao_types,                       ONLY: pao_env_type,&
      15              :                                               training_matrix_type
      16              :    USE parallel_rng_types,              ONLY: rng_stream_type
      17              : #include "./base/base_uses.f90"
      18              : 
      19              :    IMPLICIT NONE
      20              : 
      21              :    PRIVATE
      22              : 
      23              :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_ml_neuralnet'
      24              : 
      25              :    PUBLIC ::pao_ml_nn_train, pao_ml_nn_predict, pao_ml_nn_gradient
      26              : 
      27              :    ! TODO turn these into input parameters
      28              :    REAL(dp), PARAMETER   :: step_size = 0.001_dp
      29              :    INTEGER, PARAMETER    :: nlayers = 3
      30              :    REAL(dp), PARAMETER   :: convergence_eps = 1e-7_dp
      31              :    INTEGER, PARAMETER    :: max_training_cycles = 50000
      32              : 
      33              : CONTAINS
      34              : 
      35              : ! **************************************************************************************************
      36              : !> \brief Uses neural network to make a prediction
      37              : !> \param pao ...
      38              : !> \param ikind ...
      39              : !> \param descriptor ...
      40              : !> \param output ...
      41              : !> \param variance ...
      42              : ! **************************************************************************************************
      43           28 :    SUBROUTINE pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
      44              :       TYPE(pao_env_type), POINTER                        :: pao
      45              :       INTEGER, INTENT(IN)                                :: ikind
      46              :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: descriptor
      47              :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: output
      48              :       REAL(dp), INTENT(OUT)                              :: variance
      49              : 
      50              :       TYPE(training_matrix_type), POINTER                :: training_matrix
      51              : 
      52           28 :       training_matrix => pao%ml_training_matrices(ikind)
      53              : 
      54           28 :       CALL nn_eval(training_matrix%NN, input=descriptor, prediction=output)
      55              : 
      56           28 :       variance = 0.0_dp ! Neural Networks don't provide a variance
      57           28 :    END SUBROUTINE pao_ml_nn_predict
      58              : 
      59              : ! **************************************************************************************************
      60              : !> \brief Calculate gradient of neural network
      61              : !> \param pao ...
      62              : !> \param ikind ...
      63              : !> \param descriptor ...
      64              : !> \param outer_deriv ...
      65              : !> \param gradient ...
      66              : ! **************************************************************************************************
      67            4 :    SUBROUTINE pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
      68              :       TYPE(pao_env_type), POINTER                        :: pao
      69              :       INTEGER, INTENT(IN)                                :: ikind
      70              :       REAL(dp), DIMENSION(:), INTENT(IN), TARGET         :: descriptor
      71              :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: outer_deriv
      72              :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: gradient
      73              : 
      74              :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
      75              :                                                             width_out
      76            4 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: backward, forward
      77            4 :       REAL(dp), DIMENSION(:, :, :), POINTER              :: A
      78              : 
      79            4 :       A => pao%ml_training_matrices(ikind)%NN
      80              : 
      81            4 :       nlayers = SIZE(A, 1)
      82            0 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
      83            4 :       width_in = SIZE(descriptor)
      84            4 :       width_out = SIZE(outer_deriv)
      85              : 
      86           24 :       ALLOCATE (forward(0:nlayers, width), backward(0:nlayers, width))
      87              : 
      88          144 :       forward = 0.0_dp
      89           16 :       forward(0, 1:width_in) = descriptor
      90              : 
      91           16 :       DO ilayer = 1, nlayers
      92          100 :       DO i = 1, width
      93          684 :       DO j = 1, width
      94          672 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
      95              :       END DO
      96              :       END DO
      97              :       END DO
      98              : 
      99              :       ! Turning Point ------------------------------------------------------------------------------
     100          144 :       backward = 0.0_dp
     101           32 :       backward(nlayers, 1:width_out) = outer_deriv(:)
     102              : 
     103           16 :       DO ilayer = nlayers, 1, -1
     104          100 :       DO i = 1, width
     105          684 :       DO j = 1, width
     106          672 :   backward(ilayer - 1, j) = backward(ilayer - 1, j) + backward(ilayer, i)*A(ilayer, i, j)*(1.0_dp - TANH(forward(ilayer - 1, j))**2)
     107              :       END DO
     108              :       END DO
     109              :       END DO
     110              : 
     111           16 :       gradient(:) = backward(0, 1:width_in)
     112              : 
     113            4 :       DEALLOCATE (forward, backward)
     114            4 :    END SUBROUTINE pao_ml_nn_gradient
     115              : 
     116              : ! **************************************************************************************************
     117              : !> \brief Trains the neural network on given training points
     118              : !> \param pao ...
     119              : ! **************************************************************************************************
     120            4 :    SUBROUTINE pao_ml_nn_train(pao)
     121              :       TYPE(pao_env_type), POINTER                        :: pao
     122              : 
     123              :       INTEGER                                            :: i, icycle, ikind, ilayer, ipoint, j, &
     124              :                                                             npoints, width, width_in, width_out
     125              :       REAL(dp)                                           :: bak, eps, error, error1, error2, num_grad
     126            4 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: prediction
     127            4 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :, :)          :: gradient
     128              :       TYPE(rng_stream_type)                              :: rng_stream
     129              :       TYPE(training_matrix_type), POINTER                :: training_matrix
     130              : 
     131              :       ! TODO this could be parallelized over ranks
     132            8 :       DO ikind = 1, SIZE(pao%ml_training_matrices)
     133            4 :          training_matrix => pao%ml_training_matrices(ikind)
     134              : 
     135            4 :          npoints = SIZE(training_matrix%inputs, 2) ! number of points
     136            4 :          CPASSERT(SIZE(training_matrix%outputs, 2) == npoints)
     137            4 :          IF (npoints == 0) CYCLE
     138              : 
     139              :          !TODO proper output
     140            6 :          IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Training neural network for kind: ", &
     141            4 :             TRIM(training_matrix%kindname), " from ", npoints, "training points."
     142              : 
     143              :          ! determine network width and allocate it
     144            4 :          width_in = SIZE(training_matrix%inputs, 1)
     145            4 :          width_out = SIZE(training_matrix%outputs, 1)
     146            4 :          width = MAX(width_in, width_out)
     147           16 :          ALLOCATE (training_matrix%NN(nlayers, width, width))
     148              : 
     149              :          ! initialize network with random numbers from -1.0 ... +1.0
     150            4 :          rng_stream = rng_stream_type(name="pao_nn")
     151           16 :          DO ilayer = 1, nlayers
     152          100 :          DO i = 1, width
     153          684 :          DO j = 1, width
     154          672 :             training_matrix%NN(ilayer, i, j) = -1.0_dp + 2.0_dp*rng_stream%next()
     155              :          END DO
     156              :          END DO
     157              :          END DO
     158              : 
     159              :          ! train the network using backpropagation
     160           12 :          ALLOCATE (gradient(nlayers, width, width))
     161       183452 :          DO icycle = 1, max_training_cycles
     162       183450 :             error = 0.0_dp
     163     37423800 :             gradient = 0.0_dp
     164       917250 :             DO ipoint = 1, npoints
     165              :                CALL nn_backpropagate(training_matrix%NN, &
     166              :                                      input=training_matrix%inputs(:, ipoint), &
     167              :                                      goal=training_matrix%outputs(:, ipoint), &
     168              :                                      gradient=gradient, &
     169       917250 :                                      error=error)
     170              :             END DO
     171     37423800 :             training_matrix%NN(:, :, :) = training_matrix%NN - step_size*gradient
     172              : 
     173       183450 :             IF (pao%iw > 0 .AND. MOD(icycle, 100) == 0) WRITE (pao%iw, *) &
     174          917 :                "PAO|ML| ", TRIM(training_matrix%kindname), &
     175       187985 :                " training-cycle:", icycle, "SQRT(error):", SQRT(error), "grad:", SUM(gradient**2)
     176              : 
     177     37423802 :             IF (SUM(gradient**2) < convergence_eps) EXIT
     178              :          END DO
     179              : 
     180              :          ! numeric gradient for debugging ----------------------------------------------------------
     181              :          IF (.FALSE.) THEN
     182              :             eps = 1e-4_dp
     183              :             ilayer = 1
     184              :             ipoint = 1
     185              :             error = 0.0_dp
     186              :             gradient = 0.0_dp
     187              :             CALL nn_backpropagate(training_matrix%NN, &
     188              :                                   input=training_matrix%inputs(:, ipoint), &
     189              :                                   goal=training_matrix%outputs(:, ipoint), &
     190              :                                   gradient=gradient, &
     191              :                                   error=error)
     192              : 
     193              :             ALLOCATE (prediction(width_out))
     194              :             DO i = 1, width
     195              :             DO j = 1, width
     196              :                bak = training_matrix%NN(ilayer, i, j)
     197              : 
     198              :                training_matrix%NN(ilayer, i, j) = bak + eps
     199              :                CALL nn_eval(training_matrix%NN, &
     200              :                             input=training_matrix%inputs(:, ipoint), &
     201              :                             prediction=prediction)
     202              :                error1 = SUM((training_matrix%outputs(:, ipoint) - prediction)**2)
     203              : 
     204              :                training_matrix%NN(ilayer, i, j) = bak - eps
     205              :                CALL nn_eval(training_matrix%NN, &
     206              :                             input=training_matrix%inputs(:, ipoint), &
     207              :                             prediction=prediction)
     208              :                error2 = SUM((training_matrix%outputs(:, ipoint) - prediction)**2)
     209              : 
     210              :                training_matrix%NN(ilayer, i, j) = bak
     211              :                num_grad = (error1 - error2)/(2.0_dp*eps)
     212              :                IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Numeric gradient:", i, j, gradient(ilayer, i, j), num_grad
     213              : 
     214              :             END DO
     215              :             END DO
     216              :             DEALLOCATE (prediction)
     217              :          END IF
     218              :          !------------------------------------------------------------------------------------------
     219              : 
     220            4 :          DEALLOCATE (gradient)
     221              : 
     222              :          ! test training points individually
     223           12 :          ALLOCATE (prediction(width_out))
     224           20 :          DO ipoint = 1, npoints
     225              :             CALL nn_eval(training_matrix%NN, &
     226              :                          input=training_matrix%inputs(:, ipoint), &
     227           16 :                          prediction=prediction)
     228          144 :             error = MAXVAL(ABS(training_matrix%outputs(:, ipoint) - prediction))
     229           24 :             IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| ", TRIM(training_matrix%kindname), &
     230           20 :                " verify training-point:", ipoint, "SQRT(error):", SQRT(error)
     231              :          END DO
     232            8 :          DEALLOCATE (prediction)
     233              : 
     234              :       END DO
     235              : 
     236          104 :    END SUBROUTINE pao_ml_nn_train
     237              : 
     238              : ! **************************************************************************************************
     239              : !> \brief Evaluates the neural network for a given input
     240              : !> \param A ...
     241              : !> \param input ...
     242              : !> \param prediction ...
     243              : ! **************************************************************************************************
     244           44 :    SUBROUTINE nn_eval(A, input, prediction)
     245              :       REAL(dp), DIMENSION(:, :, :), INTENT(IN)           :: A
     246              :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: input
     247              :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: prediction
     248              : 
     249              :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
     250              :                                                             width_out
     251           44 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: forward
     252              : 
     253           44 :       nlayers = SIZE(A, 1)
     254           44 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
     255           44 :       width_in = SIZE(input)
     256           44 :       width_out = SIZE(prediction)
     257              : 
     258          176 :       ALLOCATE (forward(0:nlayers, width))
     259              : 
     260         1584 :       forward = 0.0_dp
     261          128 :       forward(0, 1:width_in) = input(:)
     262              : 
     263          176 :       DO ilayer = 1, nlayers
     264         1100 :       DO i = 1, width
     265         7524 :       DO j = 1, width
     266         7392 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
     267              :       END DO
     268              :       END DO
     269              :       END DO
     270              : 
     271          352 :       prediction(:) = forward(nlayers, 1:width_out)
     272              : 
     273           44 :    END SUBROUTINE nn_eval
     274              : 
     275              : ! **************************************************************************************************
     276              : !> \brief Uses backpropagation to calculate the gradient for a given training point
     277              : !> \param A ...
     278              : !> \param input ...
     279              : !> \param goal ...
     280              : !> \param error ...
     281              : !> \param gradient ...
     282              : ! **************************************************************************************************
     283       733800 :    SUBROUTINE nn_backpropagate(A, input, goal, error, gradient)
     284              :       REAL(dp), DIMENSION(:, :, :), INTENT(IN)           :: A
     285              :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: input, goal
     286              :       REAL(dp), INTENT(INOUT)                            :: error
     287              :       REAL(dp), DIMENSION(:, :, :), INTENT(INOUT)        :: gradient
     288              : 
     289              :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
     290              :                                                             width_out
     291       733800 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: prediction
     292       733800 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: backward, forward
     293              : 
     294       733800 :       nlayers = SIZE(A, 1)
     295       733800 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
     296       733800 :       width_in = SIZE(input)
     297       733800 :       width_out = SIZE(goal)
     298              : 
     299      5870400 :       ALLOCATE (forward(0:nlayers, width), prediction(width_out), backward(0:nlayers, width))
     300              : 
     301     26416800 :       forward = 0.0_dp
     302      2802800 :       forward(0, 1:width_in) = input
     303              : 
     304      2935200 :       DO ilayer = 1, nlayers
     305     18345000 :       DO i = 1, width
     306    125479800 :       DO j = 1, width
     307    123278400 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
     308              :       END DO
     309              :       END DO
     310              :       END DO
     311              : 
     312      5870400 :       prediction(:) = forward(nlayers, 1:width_out)
     313              : 
     314      5870400 :       error = error + SUM((prediction - goal)**2)
     315              : 
     316              :       ! Turning Point ------------------------------------------------------------------------------
     317     26416800 :       backward = 0.0_dp
     318      5870400 :       backward(nlayers, 1:width_out) = prediction - goal
     319              : 
     320      2935200 :       DO ilayer = nlayers, 1, -1
     321     18345000 :       DO i = 1, width
     322    125479800 :       DO j = 1, width
     323    107868600 :          gradient(ilayer, i, j) = gradient(ilayer, i, j) + 2.0_dp*backward(ilayer, i)*TANH(forward(ilayer - 1, j))
     324    123278400 :   backward(ilayer - 1, j) = backward(ilayer - 1, j) + backward(ilayer, i)*A(ilayer, i, j)*(1.0_dp - TANH(forward(ilayer - 1, j))**2)
     325              :       END DO
     326              :       END DO
     327              :       END DO
     328              : 
     329       733800 :       DEALLOCATE (forward, backward, prediction)
     330       733800 :    END SUBROUTINE nn_backpropagate
     331              : 
     332              : END MODULE pao_ml_neuralnet
        

Generated by: LCOV version 2.0-1