Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 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_layers = -1
69 : REAL(dp) :: cutoff = 0.0_dp
70 : INTEGER, DIMENSION(:), ALLOCATABLE :: kinds_mapping
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 : LOGICAL :: ignore_convergence_failure = .FALSE.
172 : CHARACTER(LEN=default_path_length) :: preopt_dm_file = ""
173 : CHARACTER(LEN=default_path_length) :: restart_file = ""
174 : TYPE(filename_type), DIMENSION(:), ALLOCATABLE :: ml_training_set
175 :
176 : INTEGER :: ml_method = -1
177 : INTEGER :: ml_prior = -1
178 : INTEGER :: ml_descriptor = -1
179 : REAL(KIND=dp) :: ml_tolerance = 0.0_dp
180 : REAL(KIND=dp) :: gp_noise_var = 0.0_dp
181 : REAL(KIND=dp) :: gp_scale = 0.0_dp
182 :
183 : ! output units
184 : INTEGER :: iw = -1
185 : INTEGER :: iw_atoms = -1
186 : INTEGER :: iw_gap = -1
187 : INTEGER :: iw_fockev = -1
188 : INTEGER :: iw_opt = -1
189 : INTEGER :: iw_mlvar = -1
190 : INTEGER :: iw_mldata = -1
191 :
192 : ! state variable
193 : INTEGER :: istep = -1
194 : REAL(KIND=dp) :: energy_prev = 0.0_dp
195 : REAL(KIND=dp) :: step_start_time = 0.0_dp
196 : REAL(KIND=dp) :: norm_G = 0.0_dp
197 : TYPE(linesearch_type) :: linesearch = linesearch_type()
198 : LOGICAL :: matrix_X_ready = .FALSE.
199 : LOGICAL :: matrix_P_ready = .FALSE.
200 : LOGICAL :: constants_ready = .FALSE.
201 : LOGICAL :: need_initial_scf = .FALSE.
202 :
203 : ! matrices
204 : TYPE(dbcsr_type) :: matrix_X
205 : TYPE(dbcsr_type) :: matrix_U0
206 : TYPE(dbcsr_type) :: matrix_H0
207 : TYPE(dbcsr_type) :: matrix_Y
208 : TYPE(dbcsr_type) :: matrix_N
209 : TYPE(dbcsr_type) :: matrix_N_diag
210 : TYPE(dbcsr_type) :: matrix_N_inv
211 : TYPE(dbcsr_type) :: matrix_N_inv_diag
212 : TYPE(dbcsr_type) :: matrix_X_orig
213 : TYPE(dbcsr_type) :: matrix_G
214 : TYPE(dbcsr_type) :: matrix_G_prev
215 : TYPE(dbcsr_type) :: matrix_D
216 : TYPE(dbcsr_type) :: matrix_D_preconed
217 : TYPE(dbcsr_type) :: matrix_V_terms
218 : TYPE(dbcsr_type) :: matrix_BFGS
219 : TYPE(dbcsr_type) :: matrix_precon
220 : TYPE(dbcsr_type) :: matrix_precon_inv
221 : TYPE(dbcsr_type) :: matrix_R
222 :
223 : TYPE(training_matrix_type), ALLOCATABLE, &
224 : DIMENSION(:) :: ml_training_matrices
225 :
226 : TYPE(pao_model_type), ALLOCATABLE, &
227 : DIMENSION(:) :: models
228 :
229 : TYPE(dbcsr_distribution_type) :: diag_distribution
230 : END TYPE pao_env_type
231 :
232 : CONTAINS
233 :
234 : ! **************************************************************************************************
235 : !> \brief Finalize the PAO environment
236 : !> \param pao ...
237 : ! **************************************************************************************************
238 98 : SUBROUTINE pao_finalize(pao)
239 : TYPE(pao_env_type) :: pao
240 :
241 : CHARACTER(len=*), PARAMETER :: routineN = 'pao_finalize'
242 :
243 : INTEGER :: handle, i
244 :
245 98 : CALL timeset(routineN, handle)
246 :
247 98 : CALL dbcsr_release(pao%matrix_X)
248 98 : CALL dbcsr_release(pao%matrix_Y)
249 98 : CALL dbcsr_release(pao%matrix_N)
250 98 : CALL dbcsr_release(pao%matrix_N_diag)
251 98 : CALL dbcsr_release(pao%matrix_N_inv)
252 98 : CALL dbcsr_release(pao%matrix_N_inv_diag)
253 98 : CALL dbcsr_release(pao%matrix_H0)
254 :
255 98 : DEALLOCATE (pao%ml_training_set)
256 98 : IF (ALLOCATED(pao%ml_training_matrices)) &
257 38 : DEALLOCATE (pao%ml_training_matrices)
258 :
259 98 : IF (ALLOCATED(pao%models)) THEN
260 12 : DO i = 1, SIZE(pao%models)
261 8 : IF (pao%models(i)%version > 0) THEN
262 8 : CALL torch_model_release(pao%models(i)%torch_model)
263 : END IF
264 12 : CALL omp_destroy_lock(pao%models(i)%lock)
265 : END DO
266 12 : DEALLOCATE (pao%models)
267 : END IF
268 :
269 98 : CALL dbcsr_distribution_release(pao%diag_distribution)
270 :
271 : !TODO: should finish printkey
272 :
273 98 : CALL timestop(handle)
274 :
275 98 : END SUBROUTINE pao_finalize
276 :
277 0 : END MODULE pao_types
|
