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