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 Main module for PAO Machine Learning
10 : !> \author Ole Schuett
11 : ! **************************************************************************************************
12 : MODULE pao_ml
13 : USE atomic_kind_types, ONLY: atomic_kind_type,&
14 : get_atomic_kind
15 : USE basis_set_types, ONLY: gto_basis_set_type
16 : USE cell_methods, ONLY: cell_create
17 : USE cell_types, ONLY: cell_type
18 : USE cp_dbcsr_api, ONLY: dbcsr_iterator_blocks_left,&
19 : dbcsr_iterator_next_block,&
20 : dbcsr_iterator_start,&
21 : dbcsr_iterator_stop,&
22 : dbcsr_iterator_type,&
23 : dbcsr_type
24 : USE kinds, ONLY: default_path_length,&
25 : default_string_length,&
26 : dp
27 : USE machine, ONLY: m_flush
28 : USE message_passing, ONLY: mp_para_env_type
29 : USE pao_input, ONLY: id2str,&
30 : pao_ml_gp,&
31 : pao_ml_lazy,&
32 : pao_ml_nn,&
33 : pao_ml_prior_mean,&
34 : pao_ml_prior_zero,&
35 : pao_rotinv_param
36 : USE pao_io, ONLY: pao_ioblock_type,&
37 : pao_iokind_type,&
38 : pao_kinds_ensure_equal,&
39 : pao_read_raw
40 : USE pao_ml_descriptor, ONLY: pao_ml_calc_descriptor
41 : USE pao_ml_gaussprocess, ONLY: pao_ml_gp_gradient,&
42 : pao_ml_gp_predict,&
43 : pao_ml_gp_train
44 : USE pao_ml_neuralnet, ONLY: pao_ml_nn_gradient,&
45 : pao_ml_nn_predict,&
46 : pao_ml_nn_train
47 : USE pao_types, ONLY: pao_env_type,&
48 : training_matrix_type
49 : USE particle_types, ONLY: particle_type
50 : USE qs_environment_types, ONLY: get_qs_env,&
51 : qs_environment_type
52 : USE qs_kind_types, ONLY: get_qs_kind,&
53 : qs_kind_type
54 : #include "./base/base_uses.f90"
55 :
56 : IMPLICIT NONE
57 :
58 : PRIVATE
59 :
60 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_ml'
61 :
62 : PUBLIC :: pao_ml_init, pao_ml_predict, pao_ml_forces
63 :
64 : ! linked list used to group training points by kind
65 : TYPE training_point_type
66 : TYPE(training_point_type), POINTER :: next => Null()
67 : REAL(dp), DIMENSION(:), ALLOCATABLE :: input
68 : REAL(dp), DIMENSION(:), ALLOCATABLE :: output
69 : END TYPE training_point_type
70 :
71 : TYPE training_list_type
72 : CHARACTER(LEN=default_string_length) :: kindname = ""
73 : TYPE(training_point_type), POINTER :: head => Null()
74 : INTEGER :: npoints = 0
75 : END TYPE training_list_type
76 :
77 : CONTAINS
78 :
79 : ! **************************************************************************************************
80 : !> \brief Initializes the learning machinery
81 : !> \param pao ...
82 : !> \param qs_env ...
83 : ! **************************************************************************************************
84 98 : SUBROUTINE pao_ml_init(pao, qs_env)
85 : TYPE(pao_env_type), POINTER :: pao
86 : TYPE(qs_environment_type), POINTER :: qs_env
87 :
88 : INTEGER :: i
89 98 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
90 : TYPE(mp_para_env_type), POINTER :: para_env
91 : TYPE(training_list_type), ALLOCATABLE, &
92 98 : DIMENSION(:) :: training_lists
93 :
94 98 : IF (SIZE(pao%ml_training_set) == 0) RETURN
95 :
96 18 : IF (pao%iw > 0) WRITE (pao%iw, *) 'PAO|ML| Initializing maschine learning...'
97 :
98 18 : IF (pao%parameterization /= pao_rotinv_param) THEN
99 0 : CPABORT("PAO maschine learning requires ROTINV parametrization")
100 : END IF
101 :
102 18 : CALL get_qs_env(qs_env, para_env=para_env, atomic_kind_set=atomic_kind_set)
103 :
104 : ! create training-set data-structure
105 74 : ALLOCATE (training_lists(SIZE(atomic_kind_set)))
106 38 : DO i = 1, SIZE(training_lists)
107 38 : CALL get_atomic_kind(atomic_kind_set(i), name=training_lists(i)%kindname)
108 : END DO
109 :
110 : ! parses training files, calculates descriptors and stores all training-points as linked lists
111 52 : DO i = 1, SIZE(pao%ml_training_set)
112 52 : CALL add_to_training_list(pao, qs_env, training_lists, filename=pao%ml_training_set(i)%fn)
113 : END DO
114 :
115 : ! ensure there there are training points for all kinds that use pao
116 18 : CALL sanity_check(qs_env, training_lists)
117 :
118 : ! turns linked lists into matrices and syncs them across ranks
119 18 : CALL training_list2matrix(training_lists, pao%ml_training_matrices, para_env)
120 :
121 : ! calculate and subtract prior
122 18 : CALL pao_ml_substract_prior(pao%ml_prior, pao%ml_training_matrices)
123 :
124 : ! print some statistics about the training set and dump it upon request
125 18 : CALL pao_ml_print(pao, pao%ml_training_matrices)
126 :
127 : ! use training-set to train model
128 18 : CALL pao_ml_train(pao)
129 :
130 116 : END SUBROUTINE pao_ml_init
131 :
132 : ! **************************************************************************************************
133 : !> \brief Reads the given file and adds its training points to linked lists.
134 : !> \param pao ...
135 : !> \param qs_env ...
136 : !> \param training_lists ...
137 : !> \param filename ...
138 : ! **************************************************************************************************
139 34 : SUBROUTINE add_to_training_list(pao, qs_env, training_lists, filename)
140 : TYPE(pao_env_type), POINTER :: pao
141 : TYPE(qs_environment_type), POINTER :: qs_env
142 : TYPE(training_list_type), DIMENSION(:) :: training_lists
143 : CHARACTER(LEN=default_path_length) :: filename
144 :
145 : CHARACTER(LEN=default_string_length) :: param
146 : INTEGER :: iatom, ikind, natoms, nkinds, nparams
147 34 : INTEGER, ALLOCATABLE, DIMENSION(:) :: atom2kind, kindsmap
148 : INTEGER, DIMENSION(2) :: ml_range
149 34 : REAL(dp), ALLOCATABLE, DIMENSION(:, :) :: hmat, positions
150 34 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
151 : TYPE(cell_type), POINTER :: cell
152 : TYPE(mp_para_env_type), POINTER :: para_env
153 34 : TYPE(pao_ioblock_type), ALLOCATABLE, DIMENSION(:) :: xblocks
154 34 : TYPE(pao_iokind_type), ALLOCATABLE, DIMENSION(:) :: kinds
155 34 : TYPE(particle_type), DIMENSION(:), POINTER :: my_particle_set
156 34 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
157 : TYPE(training_point_type), POINTER :: new_point
158 :
159 34 : NULLIFY (new_point, cell)
160 :
161 17 : IF (pao%iw > 0) WRITE (pao%iw, '(A,A)') " PAO|ML| Reading training frame from file: ", TRIM(filename)
162 :
163 34 : CALL get_qs_env(qs_env, para_env=para_env)
164 :
165 : ! parse training data on first rank
166 34 : IF (para_env%is_source()) THEN
167 17 : CALL pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks, ml_range)
168 :
169 : ! check parametrization
170 17 : IF (TRIM(param) /= TRIM(ADJUSTL(id2str(pao%parameterization)))) THEN
171 17 : CPABORT("Restart PAO parametrization does not match")
172 : END IF
173 :
174 : ! map read-in kinds onto kinds of this run
175 17 : CALL match_kinds(pao, qs_env, kinds, kindsmap)
176 17 : nkinds = SIZE(kindsmap)
177 17 : natoms = SIZE(positions, 1)
178 : END IF
179 :
180 : ! broadcast parsed raw training data
181 34 : CALL para_env%bcast(nkinds)
182 34 : CALL para_env%bcast(natoms)
183 34 : IF (.NOT. para_env%is_source()) THEN
184 17 : ALLOCATE (hmat(3, 3))
185 51 : ALLOCATE (kindsmap(nkinds))
186 51 : ALLOCATE (positions(natoms, 3))
187 51 : ALLOCATE (atom2kind(natoms))
188 : END IF
189 34 : CALL para_env%bcast(hmat)
190 34 : CALL para_env%bcast(kindsmap)
191 34 : CALL para_env%bcast(atom2kind)
192 34 : CALL para_env%bcast(positions)
193 34 : CALL para_env%bcast(ml_range)
194 :
195 34 : IF (ml_range(1) /= 1 .OR. ml_range(2) /= natoms) THEN
196 0 : CPWARN("Skipping some atoms for PAO-ML training.")
197 : END IF
198 :
199 : ! create cell from read-in h-matrix
200 34 : CALL cell_create(cell, hmat)
201 :
202 : ! create a particle_set based on read-in positions and refere to kinds of this run
203 34 : CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set, qs_kind_set=qs_kind_set)
204 476 : ALLOCATE (my_particle_set(natoms))
205 102 : DO iatom = 1, natoms
206 68 : ikind = kindsmap(atom2kind(iatom))
207 68 : my_particle_set(iatom)%atomic_kind => atomic_kind_set(ikind)
208 306 : my_particle_set(iatom)%r = positions(iatom, :)
209 : END DO
210 :
211 : ! fill linked list with training points
212 : ! Afterwards all ranks will have lists with the same number of entries,
213 : ! however the input and output arrays will only be allocated on one rank per entry.
214 : ! We farm out the expensive calculation of the descriptor across ranks.
215 102 : DO iatom = 1, natoms
216 68 : IF (iatom < ml_range(1) .OR. ml_range(2) < iatom) CYCLE
217 68 : ALLOCATE (new_point)
218 :
219 : ! training-point input, calculate descriptor only on one rank
220 68 : IF (MOD(iatom - 1, para_env%num_pe) == para_env%mepos) THEN
221 : CALL pao_ml_calc_descriptor(pao, &
222 : my_particle_set, &
223 : qs_kind_set, &
224 : cell, &
225 : iatom=iatom, &
226 34 : descriptor=new_point%input)
227 : END IF
228 :
229 : ! copy training-point output on first rank
230 68 : IF (para_env%is_source()) THEN
231 34 : nparams = SIZE(xblocks(iatom)%p, 1)
232 102 : ALLOCATE (new_point%output(nparams))
233 272 : new_point%output(:) = xblocks(iatom)%p(:, 1)
234 : END IF
235 :
236 : ! add to linked list
237 68 : ikind = kindsmap(atom2kind(iatom))
238 68 : training_lists(ikind)%npoints = training_lists(ikind)%npoints + 1
239 68 : new_point%next => training_lists(ikind)%head
240 102 : training_lists(ikind)%head => new_point
241 : END DO
242 :
243 34 : DEALLOCATE (cell, my_particle_set, hmat, kindsmap, positions, atom2kind)
244 :
245 119 : END SUBROUTINE add_to_training_list
246 :
247 : ! **************************************************************************************************
248 : !> \brief Make read-in kinds on to atomic-kinds of this run
249 : !> \param pao ...
250 : !> \param qs_env ...
251 : !> \param kinds ...
252 : !> \param kindsmap ...
253 : ! **************************************************************************************************
254 17 : SUBROUTINE match_kinds(pao, qs_env, kinds, kindsmap)
255 : TYPE(pao_env_type), POINTER :: pao
256 : TYPE(qs_environment_type), POINTER :: qs_env
257 : TYPE(pao_iokind_type), DIMENSION(:) :: kinds
258 : INTEGER, ALLOCATABLE, DIMENSION(:) :: kindsmap
259 :
260 : CHARACTER(LEN=default_string_length) :: name
261 : INTEGER :: ikind, jkind
262 17 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
263 :
264 17 : CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set)
265 :
266 17 : CPASSERT(.NOT. ALLOCATED(kindsmap))
267 51 : ALLOCATE (kindsmap(SIZE(kinds)))
268 34 : kindsmap(:) = -1
269 :
270 34 : DO ikind = 1, SIZE(kinds)
271 35 : DO jkind = 1, SIZE(atomic_kind_set)
272 18 : CALL get_atomic_kind(atomic_kind_set(jkind), name=name)
273 : ! match kinds via their name
274 18 : IF (TRIM(kinds(ikind)%name) == TRIM(name)) THEN
275 17 : CALL pao_kinds_ensure_equal(pao, qs_env, jkind, kinds(ikind))
276 17 : kindsmap(ikind) = jkind
277 17 : EXIT
278 : END IF
279 : END DO
280 : END DO
281 :
282 34 : IF (ANY(kindsmap < 1)) THEN
283 0 : CPABORT("PAO: Could not match all kinds from training set")
284 : END IF
285 17 : END SUBROUTINE match_kinds
286 :
287 : ! **************************************************************************************************
288 : !> \brief Checks that there is at least one training point per pao-enabled kind
289 : !> \param qs_env ...
290 : !> \param training_lists ...
291 : ! **************************************************************************************************
292 18 : SUBROUTINE sanity_check(qs_env, training_lists)
293 : TYPE(qs_environment_type), POINTER :: qs_env
294 : TYPE(training_list_type), DIMENSION(:), TARGET :: training_lists
295 :
296 : INTEGER :: ikind, pao_basis_size
297 : TYPE(gto_basis_set_type), POINTER :: basis_set
298 18 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
299 : TYPE(training_list_type), POINTER :: training_list
300 :
301 18 : CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set)
302 :
303 38 : DO ikind = 1, SIZE(training_lists)
304 20 : training_list => training_lists(ikind)
305 20 : IF (training_list%npoints > 0) CYCLE ! it's ok
306 2 : CALL get_qs_kind(qs_kind_set(ikind), basis_set=basis_set, pao_basis_size=pao_basis_size)
307 20 : IF (pao_basis_size /= basis_set%nsgf) THEN
308 : ! if this kind has pao enabled...
309 0 : CPABORT("Found no training-points for kind: "//TRIM(training_list%kindname))
310 : END IF
311 : END DO
312 :
313 18 : END SUBROUTINE sanity_check
314 :
315 : ! **************************************************************************************************
316 : !> \brief Turns the linked lists of training points into matrices
317 : !> \param training_lists ...
318 : !> \param training_matrices ...
319 : !> \param para_env ...
320 : ! **************************************************************************************************
321 18 : SUBROUTINE training_list2matrix(training_lists, training_matrices, para_env)
322 : TYPE(training_list_type), ALLOCATABLE, &
323 : DIMENSION(:), TARGET :: training_lists
324 : TYPE(training_matrix_type), ALLOCATABLE, &
325 : DIMENSION(:), TARGET :: training_matrices
326 : TYPE(mp_para_env_type), POINTER :: para_env
327 :
328 : INTEGER :: i, ikind, inp_size, ninputs, noutputs, &
329 : npoints, out_size
330 : TYPE(training_list_type), POINTER :: training_list
331 : TYPE(training_matrix_type), POINTER :: training_matrix
332 : TYPE(training_point_type), POINTER :: cur_point, prev_point
333 :
334 18 : CPASSERT(ALLOCATED(training_lists) .AND. .NOT. ALLOCATED(training_matrices))
335 :
336 74 : ALLOCATE (training_matrices(SIZE(training_lists)))
337 :
338 38 : DO ikind = 1, SIZE(training_lists)
339 20 : training_list => training_lists(ikind)
340 20 : training_matrix => training_matrices(ikind)
341 20 : training_matrix%kindname = training_list%kindname ! copy kindname
342 20 : npoints = training_list%npoints ! number of points
343 20 : IF (npoints == 0) THEN
344 2 : ALLOCATE (training_matrix%inputs(0, 0))
345 2 : ALLOCATE (training_matrix%outputs(0, 0))
346 2 : CYCLE
347 : END IF
348 :
349 : ! figure out size of input and output
350 18 : inp_size = 0; out_size = 0
351 18 : IF (ALLOCATED(training_list%head%input)) THEN
352 9 : inp_size = SIZE(training_list%head%input)
353 : END IF
354 18 : IF (ALLOCATED(training_list%head%output)) THEN
355 9 : out_size = SIZE(training_list%head%output)
356 : END IF
357 18 : CALL para_env%sum(inp_size)
358 18 : CALL para_env%sum(out_size)
359 :
360 : ! allocate matices to hold all training points
361 72 : ALLOCATE (training_matrix%inputs(inp_size, npoints))
362 72 : ALLOCATE (training_matrix%outputs(out_size, npoints))
363 3258 : training_matrix%inputs(:, :) = 0.0_dp
364 562 : training_matrix%outputs(:, :) = 0.0_dp
365 :
366 : ! loop over all training points, consume linked-list in the process
367 18 : ninputs = 0; noutputs = 0
368 18 : cur_point => training_list%head
369 18 : NULLIFY (training_list%head)
370 86 : DO i = 1, npoints
371 68 : IF (ALLOCATED(cur_point%input)) THEN
372 3240 : training_matrix%inputs(:, i) = cur_point%input(:)
373 34 : ninputs = ninputs + 1
374 : END IF
375 68 : IF (ALLOCATED(cur_point%output)) THEN
376 544 : training_matrix%outputs(:, i) = cur_point%output(:)
377 34 : noutputs = noutputs + 1
378 : END IF
379 : ! advance to next entry and deallocate the current one
380 68 : prev_point => cur_point
381 68 : cur_point => cur_point%next
382 86 : DEALLOCATE (prev_point)
383 : END DO
384 18 : training_list%npoints = 0 ! list is now empty
385 :
386 : ! sync training_matrix across ranks
387 18 : CALL para_env%sum(training_matrix%inputs)
388 18 : CALL para_env%sum(training_matrix%outputs)
389 :
390 : ! sanity check
391 18 : CALL para_env%sum(noutputs)
392 18 : CALL para_env%sum(ninputs)
393 36 : CPASSERT(noutputs == npoints .AND. ninputs == npoints)
394 : END DO
395 :
396 18 : END SUBROUTINE training_list2matrix
397 :
398 : ! **************************************************************************************************
399 : !> \brief TODO
400 : !> \param ml_prior ...
401 : !> \param training_matrices ...
402 : ! **************************************************************************************************
403 18 : SUBROUTINE pao_ml_substract_prior(ml_prior, training_matrices)
404 : INTEGER, INTENT(IN) :: ml_prior
405 : TYPE(training_matrix_type), DIMENSION(:), TARGET :: training_matrices
406 :
407 : INTEGER :: i, ikind, npoints, out_size
408 : TYPE(training_matrix_type), POINTER :: training_matrix
409 :
410 38 : DO ikind = 1, SIZE(training_matrices)
411 20 : training_matrix => training_matrices(ikind)
412 20 : out_size = SIZE(training_matrix%outputs, 1)
413 20 : npoints = SIZE(training_matrix%outputs, 2)
414 20 : IF (npoints == 0) CYCLE
415 54 : ALLOCATE (training_matrix%prior(out_size))
416 :
417 : ! calculate prior
418 18 : SELECT CASE (ml_prior)
419 : CASE (pao_ml_prior_zero)
420 96 : training_matrix%prior(:) = 0.0_dp
421 : CASE (pao_ml_prior_mean)
422 188 : training_matrix%prior(:) = SUM(training_matrix%outputs, 2)/REAL(npoints, dp)
423 : CASE DEFAULT
424 18 : CPABORT("PAO: unknown prior")
425 : END SELECT
426 :
427 : ! subtract prior from all training points
428 104 : DO i = 1, npoints
429 564 : training_matrix%outputs(:, i) = training_matrix%outputs(:, i) - training_matrix%prior
430 : END DO
431 : END DO
432 :
433 18 : END SUBROUTINE pao_ml_substract_prior
434 :
435 : ! **************************************************************************************************
436 : !> \brief Print some statistics about the training set and dump it upon request
437 : !> \param pao ...
438 : !> \param training_matrices ...
439 : ! **************************************************************************************************
440 18 : SUBROUTINE pao_ml_print(pao, training_matrices)
441 : TYPE(pao_env_type), POINTER :: pao
442 : TYPE(training_matrix_type), DIMENSION(:), TARGET :: training_matrices
443 :
444 : INTEGER :: i, ikind, N, npoints
445 : TYPE(training_matrix_type), POINTER :: training_matrix
446 :
447 : ! dump training data
448 18 : IF (pao%iw_mldata > 0) THEN
449 2 : DO ikind = 1, SIZE(training_matrices)
450 1 : training_matrix => training_matrices(ikind)
451 1 : npoints = SIZE(training_matrix%outputs, 2)
452 6 : DO i = 1, npoints
453 4 : WRITE (pao%iw_mldata, *) "PAO|ML| training-point kind: ", TRIM(training_matrix%kindname), &
454 8 : " point:", i, " in:", training_matrix%inputs(:, i), &
455 37 : " out:", training_matrix%outputs(:, i)
456 : END DO
457 : END DO
458 1 : CALL m_flush(pao%iw_mldata)
459 : END IF
460 :
461 : ! print stats
462 18 : IF (pao%iw > 0) THEN
463 19 : DO ikind = 1, SIZE(training_matrices)
464 10 : training_matrix => training_matrices(ikind)
465 30 : N = SIZE(training_matrix%inputs)
466 10 : IF (N == 0) CYCLE
467 : WRITE (pao%iw, "(A,I3,A,E10.1,1X,E10.1,1X,E10.1)") " PAO|ML| Descriptor for kind: "// &
468 9 : TRIM(training_matrix%kindname)//" size: ", &
469 9 : SIZE(training_matrix%inputs, 1), " min/mean/max: ", &
470 1629 : MINVAL(training_matrix%inputs), &
471 1629 : SUM(training_matrix%inputs)/REAL(N, dp), &
472 1648 : MAXVAL(training_matrix%inputs)
473 : END DO
474 : END IF
475 :
476 18 : END SUBROUTINE pao_ml_print
477 :
478 : ! **************************************************************************************************
479 : !> \brief Calls the actual learning algorthim to traing on the given matrices
480 : !> \param pao ...
481 : ! **************************************************************************************************
482 18 : SUBROUTINE pao_ml_train(pao)
483 : TYPE(pao_env_type), POINTER :: pao
484 :
485 : CHARACTER(len=*), PARAMETER :: routineN = 'pao_ml_train'
486 :
487 : INTEGER :: handle
488 :
489 18 : CALL timeset(routineN, handle)
490 :
491 28 : SELECT CASE (pao%ml_method)
492 : CASE (pao_ml_gp)
493 10 : CALL pao_ml_gp_train(pao)
494 : CASE (pao_ml_nn)
495 4 : CALL pao_ml_nn_train(pao)
496 : CASE (pao_ml_lazy)
497 : ! nothing to do
498 : CASE DEFAULT
499 18 : CPABORT("PAO: unknown machine learning scheme")
500 : END SELECT
501 :
502 18 : CALL timestop(handle)
503 :
504 18 : END SUBROUTINE pao_ml_train
505 :
506 : ! **************************************************************************************************
507 : !> \brief Fills pao%matrix_X based on machine learning predictions
508 : !> \param pao ...
509 : !> \param qs_env ...
510 : ! **************************************************************************************************
511 138 : SUBROUTINE pao_ml_predict(pao, qs_env)
512 : TYPE(pao_env_type), POINTER :: pao
513 : TYPE(qs_environment_type), POINTER :: qs_env
514 :
515 : CHARACTER(len=*), PARAMETER :: routineN = 'pao_ml_predict'
516 :
517 : INTEGER :: acol, arow, handle, iatom, ikind, natoms
518 138 : REAL(dp), ALLOCATABLE, DIMENSION(:) :: descriptor, variances
519 138 : REAL(dp), DIMENSION(:, :), POINTER :: block_X
520 138 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
521 : TYPE(cell_type), POINTER :: cell
522 : TYPE(dbcsr_iterator_type) :: iter
523 : TYPE(mp_para_env_type), POINTER :: para_env
524 138 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
525 138 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
526 :
527 138 : CALL timeset(routineN, handle)
528 :
529 : CALL get_qs_env(qs_env, &
530 : para_env=para_env, &
531 : cell=cell, &
532 : particle_set=particle_set, &
533 : atomic_kind_set=atomic_kind_set, &
534 : qs_kind_set=qs_kind_set, &
535 138 : natom=natoms)
536 :
537 : ! fill matrix_X
538 414 : ALLOCATE (variances(natoms))
539 138 : variances(:) = 0.0_dp
540 : !$OMP PARALLEL DEFAULT(NONE) SHARED(pao,qs_env,particle_set,qs_kind_set,cell,variances) &
541 138 : !$OMP PRIVATE(iter,arow,acol,iatom,ikind,descriptor,block_X)
542 : CALL dbcsr_iterator_start(iter, pao%matrix_X)
543 : DO WHILE (dbcsr_iterator_blocks_left(iter))
544 : CALL dbcsr_iterator_next_block(iter, arow, acol, block_X)
545 : iatom = arow; CPASSERT(arow == acol)
546 : CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
547 : IF (SIZE(block_X) == 0) CYCLE ! pao disabled for iatom
548 :
549 : ! calculate descriptor
550 : CALL pao_ml_calc_descriptor(pao, &
551 : particle_set, &
552 : qs_kind_set, &
553 : cell, &
554 : iatom, &
555 : descriptor)
556 :
557 : ! call actual machine learning for prediction
558 : CALL pao_ml_predict_low(pao, ikind=ikind, &
559 : descriptor=descriptor, &
560 : output=block_X(:, 1), &
561 : variance=variances(iatom))
562 :
563 : DEALLOCATE (descriptor)
564 :
565 : !add prior
566 : block_X(:, 1) = block_X(:, 1) + pao%ml_training_matrices(ikind)%prior
567 : END DO
568 : CALL dbcsr_iterator_stop(iter)
569 : !$OMP END PARALLEL
570 :
571 : ! print variances
572 138 : CALL para_env%sum(variances)
573 138 : IF (pao%iw_mlvar > 0) THEN
574 3 : DO iatom = 1, natoms
575 3 : WRITE (pao%iw_mlvar, *) "PAO|ML| atom:", iatom, " prediction variance:", variances(iatom)
576 : END DO
577 1 : CALL m_flush(pao%iw_mlvar)
578 : END IF
579 :
580 : ! one-line summary
581 207 : IF (pao%iw > 0) WRITE (pao%iw, "(A,E20.10,A,T71,I10)") " PAO|ML| max prediction variance:", &
582 485 : MAXVAL(variances), " for atom:", MAXLOC(variances)
583 :
584 416 : IF (MAXVAL(variances) > pao%ml_tolerance) THEN
585 0 : CPABORT("Variance of prediction above ML_TOLERANCE.")
586 : END IF
587 :
588 138 : DEALLOCATE (variances)
589 :
590 138 : CALL timestop(handle)
591 :
592 276 : END SUBROUTINE pao_ml_predict
593 :
594 : ! **************************************************************************************************
595 : !> \brief Queries the actual learning algorthim to make a prediction
596 : !> \param pao ...
597 : !> \param ikind ...
598 : !> \param descriptor ...
599 : !> \param output ...
600 : !> \param variance ...
601 : ! **************************************************************************************************
602 138 : SUBROUTINE pao_ml_predict_low(pao, ikind, descriptor, output, variance)
603 : TYPE(pao_env_type), POINTER :: pao
604 : INTEGER, INTENT(IN) :: ikind
605 : REAL(dp), DIMENSION(:), INTENT(IN) :: descriptor
606 : REAL(dp), DIMENSION(:), INTENT(OUT) :: output
607 : REAL(dp), INTENT(OUT) :: variance
608 :
609 220 : SELECT CASE (pao%ml_method)
610 : CASE (pao_ml_gp)
611 82 : CALL pao_ml_gp_predict(pao, ikind, descriptor, output, variance)
612 : CASE (pao_ml_nn)
613 28 : CALL pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
614 : CASE (pao_ml_lazy)
615 224 : output = 0.0_dp ! let's be really lazy and just rely on the prior
616 28 : variance = 0
617 : CASE DEFAULT
618 138 : CPABORT("PAO: unknown machine learning scheme")
619 : END SELECT
620 :
621 138 : END SUBROUTINE pao_ml_predict_low
622 :
623 : ! **************************************************************************************************
624 : !> \brief Calculate forces contributed by machine learning
625 : !> \param pao ...
626 : !> \param qs_env ...
627 : !> \param matrix_G ...
628 : !> \param forces ...
629 : ! **************************************************************************************************
630 18 : SUBROUTINE pao_ml_forces(pao, qs_env, matrix_G, forces)
631 : TYPE(pao_env_type), POINTER :: pao
632 : TYPE(qs_environment_type), POINTER :: qs_env
633 : TYPE(dbcsr_type) :: matrix_G
634 : REAL(dp), DIMENSION(:, :), INTENT(INOUT) :: forces
635 :
636 : CHARACTER(len=*), PARAMETER :: routineN = 'pao_ml_forces'
637 :
638 : INTEGER :: acol, arow, handle, iatom, ikind
639 18 : REAL(dp), ALLOCATABLE, DIMENSION(:) :: descr_grad, descriptor
640 18 : REAL(dp), DIMENSION(:, :), POINTER :: block_G
641 18 : TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
642 : TYPE(cell_type), POINTER :: cell
643 : TYPE(dbcsr_iterator_type) :: iter
644 : TYPE(mp_para_env_type), POINTER :: para_env
645 18 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
646 18 : TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
647 :
648 18 : CALL timeset(routineN, handle)
649 :
650 : CALL get_qs_env(qs_env, &
651 : para_env=para_env, &
652 : cell=cell, &
653 : particle_set=particle_set, &
654 : atomic_kind_set=atomic_kind_set, &
655 18 : qs_kind_set=qs_kind_set)
656 :
657 : !$OMP PARALLEL DEFAULT(NONE) SHARED(pao,matrix_G,particle_set,qs_kind_set,cell) &
658 : !$OMP REDUCTION(+:forces) &
659 18 : !$OMP PRIVATE(iter,arow,acol,iatom,ikind,block_G,descriptor,descr_grad)
660 : CALL dbcsr_iterator_start(iter, matrix_G)
661 : DO WHILE (dbcsr_iterator_blocks_left(iter))
662 : CALL dbcsr_iterator_next_block(iter, arow, acol, block_G)
663 : iatom = arow; CPASSERT(arow == acol)
664 : CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
665 : IF (SIZE(block_G) == 0) CYCLE ! pao disabled for iatom
666 :
667 : ! calculate descriptor
668 : CALL pao_ml_calc_descriptor(pao, &
669 : particle_set, &
670 : qs_kind_set, &
671 : cell, &
672 : iatom=iatom, &
673 : descriptor=descriptor)
674 :
675 : ! calcaulte derivate of machine learning prediction
676 : CALL pao_ml_gradient_low(pao, ikind=ikind, &
677 : descriptor=descriptor, &
678 : outer_deriv=block_G(:, 1), &
679 : gradient=descr_grad)
680 :
681 : ! calculate force contributions from descriptor
682 : CALL pao_ml_calc_descriptor(pao, &
683 : particle_set, &
684 : qs_kind_set, &
685 : cell, &
686 : iatom=iatom, &
687 : descr_grad=descr_grad, &
688 : forces=forces)
689 :
690 : DEALLOCATE (descriptor, descr_grad)
691 : END DO
692 : CALL dbcsr_iterator_stop(iter)
693 : !$OMP END PARALLEL
694 :
695 18 : CALL timestop(handle)
696 :
697 36 : END SUBROUTINE pao_ml_forces
698 :
699 : ! **************************************************************************************************
700 : !> \brief Calculate gradient of machine learning algorithm
701 : !> \param pao ...
702 : !> \param ikind ...
703 : !> \param descriptor ...
704 : !> \param outer_deriv ...
705 : !> \param gradient ...
706 : ! **************************************************************************************************
707 18 : SUBROUTINE pao_ml_gradient_low(pao, ikind, descriptor, outer_deriv, gradient)
708 : TYPE(pao_env_type), POINTER :: pao
709 : INTEGER, INTENT(IN) :: ikind
710 : REAL(dp), DIMENSION(:), INTENT(IN) :: descriptor, outer_deriv
711 : REAL(dp), ALLOCATABLE, DIMENSION(:) :: gradient
712 :
713 54 : ALLOCATE (gradient(SIZE(descriptor)))
714 :
715 28 : SELECT CASE (pao%ml_method)
716 : CASE (pao_ml_gp)
717 10 : CALL pao_ml_gp_gradient(pao, ikind, descriptor, outer_deriv, gradient)
718 : CASE (pao_ml_nn)
719 4 : CALL pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
720 : CASE (pao_ml_lazy)
721 4 : gradient = 0.0_dp
722 : CASE DEFAULT
723 18 : CPABORT("PAO: unknown machine learning scheme")
724 : END SELECT
725 :
726 18 : END SUBROUTINE pao_ml_gradient_low
727 :
728 0 : END MODULE pao_ml
|
