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