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 Types used by the PAO machinery
      10             : !> \author Ole Schuett
      11             : ! **************************************************************************************************
      12             : MODULE pao_types
      13             :    USE cp_dbcsr_api,                    ONLY: dbcsr_distribution_release,&
      14             :                                               dbcsr_distribution_type,&
      15             :                                               dbcsr_release,&
      16             :                                               dbcsr_type
      17             :    USE kinds,                           ONLY: default_path_length,&
      18             :                                               default_string_length,&
      19             :                                               dp
      20             :    USE linesearch,                      ONLY: linesearch_type
      21             :    USE torch_api,                       ONLY: torch_model_release,&
      22             :                                               torch_model_type
      23             : #include "./base/base_uses.f90"
      24             : 
      25             :    IMPLICIT NONE
      26             : 
      27             :    PRIVATE
      28             : 
      29             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_types'
      30             : 
      31             :    PUBLIC :: pao_env_type, training_matrix_type, pao_model_type, pao_finalize
      32             : 
      33             :    TYPE filename_type
      34             :       CHARACTER(LEN=default_path_length) :: fn = ""
      35             :    END TYPE filename_type
      36             : 
      37             : ! **************************************************************************************************
      38             : !> \brief PAO machine learning data for one atomic kind
      39             : !> \var kindname                  name of atomic kind
      40             : !> \var inputs                    training points
      41             : !> \var outputs                   training points
      42             : !> \var prior                     constant prior which is added to prediction
      43             : !> \var NN                        trained neural network
      44             : !> \var GP                        trained gaussian process
      45             : ! **************************************************************************************************
      46             :    TYPE training_matrix_type
      47             :       CHARACTER(LEN=default_string_length)      :: kindname = ""
      48             :       REAL(dp), DIMENSION(:, :), ALLOCATABLE    :: inputs
      49             :       REAL(dp), DIMENSION(:, :), ALLOCATABLE    :: outputs
      50             :       REAL(dp), DIMENSION(:), ALLOCATABLE       :: prior
      51             :       REAL(dp), DIMENSION(:, :, :), ALLOCATABLE :: NN ! Neural Network
      52             :       REAL(dp), DIMENSION(:, :), ALLOCATABLE    :: GP ! Gaussian Process
      53             :    END TYPE training_matrix_type
      54             : 
      55             : ! **************************************************************************************************
      56             : !> \brief PAO-ML model for a single atomic kind.
      57             : ! **************************************************************************************************
      58             :    TYPE pao_model_type
      59             :       TYPE(torch_model_type)                    :: torch_model
      60             :       INTEGER                                   :: version = -1
      61             :       CHARACTER(LEN=default_string_length)      :: kind_name = ""
      62             :       INTEGER                                   :: atomic_number = -1
      63             :       CHARACTER(LEN=default_string_length)      :: prim_basis_name = ""
      64             :       INTEGER                                   :: prim_basis_size = -1
      65             :       INTEGER                                   :: pao_basis_size = -1
      66             :       INTEGER                                   :: num_neighbors = -1
      67             :       REAL(dp)                                  :: cutoff = 0.0_dp
      68             :       INTEGER, DIMENSION(:), ALLOCATABLE        :: feature_kinds
      69             :    END TYPE pao_model_type
      70             : 
      71             : ! **************************************************************************************************
      72             : !> \brief The PAO environment type
      73             : !> \var eps_pao                   parsed input parameter
      74             : !> \var cg_reset_limit            parsed input parameter
      75             : !> \var mixing                    parsed input parameter
      76             : !> \var regularization            parsed input parameter
      77             : !> \var penalty_dist              parsed input parameter
      78             : !> \var penalty_strength          parsed input parameter
      79             : !> \var check_unitary_tol         parsed input parameter
      80             : !> \var check_grad_tol            parsed input parameter
      81             : !> \var num_grad_eps              parsed input parameter
      82             : !> \var eps_pgf                   parsed input parameter
      83             : !> \var linpot_precon_delta       parsed input parameter
      84             : !> \var linpot_init_delta         parsed input parameter
      85             : !> \var linpot_regu_delta         parsed input parameter
      86             : !> \var linpot_regu_strength      parsed input parameter
      87             : !> \var num_grad_order            parsed input parameter
      88             : !> \var max_pao                   parsed input parameter
      89             : !> \var max_cycles                parsed input parameter
      90             : !> \var write_cycles              parsed input parameter
      91             : !> \var parameterization          parsed input parameter
      92             : !> \var optimizer                 parsed input parameter
      93             : !> \var cg_init_steps             parsed input parameter
      94             : !> \var preopt_dm_file            parsed input parameter
      95             : !> \var restart_file              parsed input parameter
      96             : !> \var ml_training_set           parsed input parameter
      97             : !> \var ml_method                 parsed input parameter
      98             : !> \var ml_prior                  parsed input parameter
      99             : !> \var ml_descriptor             parsed input parameter
     100             : !> \var ml_tolerance              parsed input parameter
     101             : !> \var gp_noise_var              parsed input parameter
     102             : !> \var gp_scale                  parsed input parameter
     103             : !> \var precondition              parsed input parameter
     104             : !> \var iw                        output unit for pao in general
     105             : !> \var iw_atoms                  output unit for one line summary for each atom
     106             : !> \var iw_gap                    output unit for gap of the fock matrix
     107             : !> \var iw_fockev                 output unit for eigenvalues of the fock matrix
     108             : !> \var iw_opt                    output unit for pao optimizer
     109             : !> \var iw_mlvar                  output unit for variances of machine learning predictions
     110             : !> \var iw_mldata                 output unit for dumping training data used for machine learning
     111             : !> \var istep                     counts pao iterations, ie. number of pao energy evaluations
     112             : !> \var energy_prev               energy of previous pao step
     113             : !> \var step_start_time           timestamp of when current pao step started
     114             : !> \var norm_G                    frobenius-norm of matrix_G or matrix_G_preconed
     115             : !> \var linesearch                holds linesearch state
     116             : !> \var matrix_X_ready            set when matrix_X is initialized
     117             : !> \var matrix_P_ready            set when density matrix is initialized
     118             : !> \var constants_ready           set when stuff, which does not depend of atomic positions is ready
     119             : !> \var need_initial_scf          set when the initial density matrix is not self-consistend
     120             : !> \var matrix_X                  parameters of pao basis, which eventually determine matrix_U. Uses diag_distribution.
     121             : !> \var matrix_U0                 constant pre-rotation which serves as initial guess for exp-parametrization. Uses diag_distribution.
     122             : !> \var matrix_H0                 Diagonal blocks of core hamiltonian, uses diag_distribution
     123             : !> \var matrix_Y                  selector matrix which translates between primary and pao basis.
     124             : !>                                basically a block diagonal "rectangular identity matrix". Uses s_matrix-distribution.
     125             : !> \var matrix_N                  diagonal matrix filled with 1/sqrt(S) from primary overlap matrix. Uses s_matrix-distribution.
     126             : !> \var matrix_N_diag             copy of matrix_N using diag_distribution
     127             : !> \var matrix_N_inv              diagonal matrix filled with sqrt(S) from primary overlap matrix. Uses s_matrix-distribution.
     128             : !> \var matrix_N_inv_diag         copy of matrix_N_inv using diag_distribution
     129             : !> \var matrix_X_orig             copy made of matrix_X at beginning of optimization cylce, used for mixing. Uses diag_distribution.
     130             : !> \var matrix_G                  derivative of pao-energy wrt to matrix_X, ie. the pao-gradient. Uses diag_distribution.
     131             : !> \var matrix_G_prev             copy of gradient from previous step, used for conjugate gradient method. Uses diag_distribution.
     132             : !> \var matrix_D                  Current line-search direction,  used for conjugate gradient method. Uses diag_distribution.
     133             : !> \var matrix_D_preconed         Current line-search direction with preconditioner applied.
     134             : !>                                This copy is keept, because application of inverse preconditioner
     135             : !>                                introduces too much numeric noise. Uses diag_distribution.
     136             : !> \var matrix_V_terms            Potential terms, used by linpot and gth parametrization, Uses diag_distribution.
     137             : !> \var matrix_BFGS               Approximate inverse hessian, used by BFGS optimizer, Uses diag_distribution.
     138             : !> \var matrix_precon             preconditioner, uses diag_distribution.
     139             : !> \var matrix_precon_inv         inverse of matrix_precon, uses diag_distribution.
     140             : !> \var matrix_R                  Rgularization, uses diag_distribution
     141             : !> \var ml_training_matrices      holds training data and trained machine learning model
     142             : !> \var diag_distribution         DBCSR distribution to spreads diagonal blocks evenly across ranks
     143             : ! **************************************************************************************************
     144             :    TYPE pao_env_type
     145             :       ! input values
     146             :       REAL(KIND=dp)                    :: eps_pao = 0.0_dp
     147             :       REAL(KIND=dp)                    :: cg_reset_limit = 0.1_dp
     148             :       REAL(KIND=dp)                    :: mixing = 0.0_dp
     149             :       REAL(KIND=dp)                    :: regularization = 0.0_dp
     150             :       REAL(KIND=dp)                    :: penalty_dist = 0.0_dp
     151             :       REAL(KIND=dp)                    :: penalty_strength = 0.0_dp
     152             :       REAL(KIND=dp)                    :: check_unitary_tol = 0.0_dp
     153             :       REAL(KIND=dp)                    :: check_grad_tol = 0.0_dp
     154             :       REAL(KIND=dp)                    :: num_grad_eps = 0.0_dp
     155             :       REAL(KIND=dp)                    :: eps_pgf = 0.0_dp
     156             :       REAL(KIND=dp)                    :: linpot_precon_delta = 0.0_dp
     157             :       REAL(KIND=dp)                    :: linpot_init_delta = 0.0_dp
     158             :       REAL(KIND=dp)                    :: linpot_regu_delta = 0.0_dp
     159             :       REAL(KIND=dp)                    :: linpot_regu_strength = 0.0_dp
     160             :       INTEGER                          :: num_grad_order = -1
     161             :       INTEGER                          :: max_pao = -1
     162             :       INTEGER                          :: max_cycles = -1
     163             :       INTEGER                          :: write_cycles = -1
     164             :       INTEGER                          :: parameterization = -1
     165             :       INTEGER                          :: optimizer = -1
     166             :       INTEGER                          :: cg_init_steps = -1
     167             :       LOGICAL                          :: precondition = .FALSE.
     168             :       CHARACTER(LEN=default_path_length) :: preopt_dm_file = ""
     169             :       CHARACTER(LEN=default_path_length) :: restart_file = ""
     170             :       TYPE(filename_type), DIMENSION(:), ALLOCATABLE :: ml_training_set
     171             : 
     172             :       INTEGER                          :: ml_method = -1
     173             :       INTEGER                          :: ml_prior = -1
     174             :       INTEGER                          :: ml_descriptor = -1
     175             :       REAL(KIND=dp)                    :: ml_tolerance = 0.0_dp
     176             :       REAL(KIND=dp)                    :: gp_noise_var = 0.0_dp
     177             :       REAL(KIND=dp)                    :: gp_scale = 0.0_dp
     178             : 
     179             :       ! output units
     180             :       INTEGER                          :: iw = -1
     181             :       INTEGER                          :: iw_atoms = -1
     182             :       INTEGER                          :: iw_gap = -1
     183             :       INTEGER                          :: iw_fockev = -1
     184             :       INTEGER                          :: iw_opt = -1
     185             :       INTEGER                          :: iw_mlvar = -1
     186             :       INTEGER                          :: iw_mldata = -1
     187             : 
     188             :       ! state variable
     189             :       INTEGER                          :: istep = -1
     190             :       REAL(KIND=dp)                    :: energy_prev = 0.0_dp
     191             :       REAL(KIND=dp)                    :: step_start_time = 0.0_dp
     192             :       REAL(KIND=dp)                    :: norm_G = 0.0_dp
     193             :       TYPE(linesearch_type)            :: linesearch = linesearch_type()
     194             :       LOGICAL                          :: matrix_X_ready = .FALSE.
     195             :       LOGICAL                          :: matrix_P_ready = .FALSE.
     196             :       LOGICAL                          :: constants_ready = .FALSE.
     197             :       LOGICAL                          :: need_initial_scf = .FALSE.
     198             : 
     199             :       ! matrices
     200             :       TYPE(dbcsr_type)              :: matrix_X
     201             :       TYPE(dbcsr_type)              :: matrix_U0
     202             :       TYPE(dbcsr_type)              :: matrix_H0
     203             :       TYPE(dbcsr_type)              :: matrix_Y
     204             :       TYPE(dbcsr_type)              :: matrix_N
     205             :       TYPE(dbcsr_type)              :: matrix_N_diag
     206             :       TYPE(dbcsr_type)              :: matrix_N_inv
     207             :       TYPE(dbcsr_type)              :: matrix_N_inv_diag
     208             :       TYPE(dbcsr_type)              :: matrix_X_orig
     209             :       TYPE(dbcsr_type)              :: matrix_G
     210             :       TYPE(dbcsr_type)              :: matrix_G_prev
     211             :       TYPE(dbcsr_type)              :: matrix_D
     212             :       TYPE(dbcsr_type)              :: matrix_D_preconed
     213             :       TYPE(dbcsr_type)              :: matrix_V_terms
     214             :       TYPE(dbcsr_type)              :: matrix_BFGS
     215             :       TYPE(dbcsr_type)              :: matrix_precon
     216             :       TYPE(dbcsr_type)              :: matrix_precon_inv
     217             :       TYPE(dbcsr_type)              :: matrix_R
     218             : 
     219             :       TYPE(training_matrix_type), ALLOCATABLE, &
     220             :          DIMENSION(:)                  :: ml_training_matrices
     221             : 
     222             :       TYPE(pao_model_type), ALLOCATABLE, &
     223             :          DIMENSION(:)                  :: models
     224             : 
     225             :       TYPE(dbcsr_distribution_type)     :: diag_distribution
     226             :    END TYPE
     227             : 
     228             : CONTAINS
     229             : 
     230             : ! **************************************************************************************************
     231             : !> \brief Finalize the PAO environment
     232             : !> \param pao ...
     233             : ! **************************************************************************************************
     234          96 :    SUBROUTINE pao_finalize(pao)
     235             :       TYPE(pao_env_type)                                 :: pao
     236             : 
     237             :       CHARACTER(len=*), PARAMETER                        :: routineN = 'pao_finalize'
     238             : 
     239             :       INTEGER                                            :: handle, i
     240             : 
     241          96 :       CALL timeset(routineN, handle)
     242             : 
     243          96 :       CALL dbcsr_release(pao%matrix_X)
     244          96 :       CALL dbcsr_release(pao%matrix_Y)
     245          96 :       CALL dbcsr_release(pao%matrix_N)
     246          96 :       CALL dbcsr_release(pao%matrix_N_diag)
     247          96 :       CALL dbcsr_release(pao%matrix_N_inv)
     248          96 :       CALL dbcsr_release(pao%matrix_N_inv_diag)
     249          96 :       CALL dbcsr_release(pao%matrix_H0)
     250             : 
     251          96 :       DEALLOCATE (pao%ml_training_set)
     252          96 :       IF (ALLOCATED(pao%ml_training_matrices)) &
     253          38 :          DEALLOCATE (pao%ml_training_matrices)
     254             : 
     255          96 :       IF (ALLOCATED(pao%models)) THEN
     256           6 :          DO i = 1, SIZE(pao%models)
     257           6 :             IF (pao%models(i)%version > 0) THEN
     258           4 :                CALL torch_model_release(pao%models(i)%torch_model)
     259             :             END IF
     260             :          END DO
     261           6 :          DEALLOCATE (pao%models)
     262             :       END IF
     263             : 
     264          96 :       CALL dbcsr_distribution_release(pao%diag_distribution)
     265             : 
     266             :       !TODO: should finish printkey
     267             : 
     268          96 :       CALL timestop(handle)
     269             : 
     270          96 :    END SUBROUTINE pao_finalize
     271             : 
     272           0 : END MODULE pao_types