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
|
