(git:b279b6b)
pao_ml.F
Go to the documentation of this file.
1 !--------------------------------------------------------------------------------------------------!
2 ! CP2K: A general program to perform molecular dynamics simulations !
3 ! Copyright 2000-2024 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 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  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  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  DIMENSION(:) :: training_lists
93 
94  IF (SIZE(pao%ml_training_set) == 0) RETURN
95 
96  IF (pao%iw > 0) WRITE (pao%iw, *) 'PAO|ML| Initializing maschine learning...'
97 
98  IF (pao%parameterization /= pao_rotinv_param) &
99  cpabort("PAO maschine learning requires ROTINV parametrization")
100 
101  CALL get_qs_env(qs_env, para_env=para_env, atomic_kind_set=atomic_kind_set)
102 
103  ! create training-set data-structure
104  ALLOCATE (training_lists(SIZE(atomic_kind_set)))
105  DO i = 1, SIZE(training_lists)
106  CALL get_atomic_kind(atomic_kind_set(i), name=training_lists(i)%kindname)
107  END DO
108 
109  ! parses training files, calculates descriptors and stores all training-points as linked lists
110  DO i = 1, SIZE(pao%ml_training_set)
111  CALL add_to_training_list(pao, qs_env, training_lists, filename=pao%ml_training_set(i)%fn)
112  END DO
113 
114  ! ensure there there are training points for all kinds that use pao
115  CALL sanity_check(qs_env, training_lists)
116 
117  ! turns linked lists into matrices and syncs them across ranks
118  CALL training_list2matrix(training_lists, pao%ml_training_matrices, para_env)
119 
120  ! calculate and subtract prior
121  CALL pao_ml_substract_prior(pao%ml_prior, pao%ml_training_matrices)
122 
123  ! print some statistics about the training set and dump it upon request
124  CALL pao_ml_print(pao, pao%ml_training_matrices)
125 
126  ! use training-set to train model
127  CALL pao_ml_train(pao)
128 
129  END SUBROUTINE pao_ml_init
130 
131 ! **************************************************************************************************
132 !> \brief Reads the given file and adds its training points to linked lists.
133 !> \param pao ...
134 !> \param qs_env ...
135 !> \param training_lists ...
136 !> \param filename ...
137 ! **************************************************************************************************
138  SUBROUTINE add_to_training_list(pao, qs_env, training_lists, filename)
139  TYPE(pao_env_type), POINTER :: pao
140  TYPE(qs_environment_type), POINTER :: qs_env
141  TYPE(training_list_type), DIMENSION(:) :: training_lists
142  CHARACTER(LEN=default_path_length) :: filename
143 
144  CHARACTER(LEN=default_string_length) :: param
145  INTEGER :: iatom, ikind, natoms, nkinds, nparams
146  INTEGER, ALLOCATABLE, DIMENSION(:) :: atom2kind, kindsmap
147  INTEGER, DIMENSION(2) :: ml_range
148  REAL(dp), ALLOCATABLE, DIMENSION(:, :) :: hmat, positions
149  TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
150  TYPE(cell_type), POINTER :: cell
151  TYPE(mp_para_env_type), POINTER :: para_env
152  TYPE(pao_ioblock_type), ALLOCATABLE, DIMENSION(:) :: xblocks
153  TYPE(pao_iokind_type), ALLOCATABLE, DIMENSION(:) :: kinds
154  TYPE(particle_type), DIMENSION(:), POINTER :: my_particle_set
155  TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
156  TYPE(training_point_type), POINTER :: new_point
157 
158  NULLIFY (new_point, cell)
159 
160  IF (pao%iw > 0) WRITE (pao%iw, '(A,A)') " PAO|ML| Reading training frame from file: ", trim(filename)
161 
162  CALL get_qs_env(qs_env, para_env=para_env)
163 
164  ! parse training data on first rank
165  IF (para_env%is_source()) THEN
166  CALL pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks, ml_range)
167 
168  ! check parametrization
169  IF (trim(param) .NE. trim(adjustl(id2str(pao%parameterization)))) &
170  cpabort("Restart PAO parametrization does not match")
171 
172  ! map read-in kinds onto kinds of this run
173  CALL match_kinds(pao, qs_env, kinds, kindsmap)
174  nkinds = SIZE(kindsmap)
175  natoms = SIZE(positions, 1)
176  END IF
177 
178  ! broadcast parsed raw training data
179  CALL para_env%bcast(nkinds)
180  CALL para_env%bcast(natoms)
181  IF (.NOT. para_env%is_source()) THEN
182  ALLOCATE (hmat(3, 3))
183  ALLOCATE (kindsmap(nkinds))
184  ALLOCATE (positions(natoms, 3))
185  ALLOCATE (atom2kind(natoms))
186  END IF
187  CALL para_env%bcast(hmat)
188  CALL para_env%bcast(kindsmap)
189  CALL para_env%bcast(atom2kind)
190  CALL para_env%bcast(positions)
191  CALL para_env%bcast(ml_range)
192 
193  IF (ml_range(1) /= 1 .OR. ml_range(2) /= natoms) &
194  cpwarn("Skipping some atoms for PAO-ML training.")
195 
196  ! create cell from read-in h-matrix
197  CALL cell_create(cell, hmat)
198 
199  ! create a particle_set based on read-in positions and refere to kinds of this run
200  CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set, qs_kind_set=qs_kind_set)
201  ALLOCATE (my_particle_set(natoms))
202  DO iatom = 1, natoms
203  ikind = kindsmap(atom2kind(iatom))
204  my_particle_set(iatom)%atomic_kind => atomic_kind_set(ikind)
205  my_particle_set(iatom)%r = positions(iatom, :)
206  END DO
207 
208  ! fill linked list with training points
209  ! Afterwards all ranks will have lists with the same number of entries,
210  ! however the input and output arrays will only be allocated on one rank per entry.
211  ! We farm out the expensive calculation of the descriptor across ranks.
212  DO iatom = 1, natoms
213  IF (iatom < ml_range(1) .OR. ml_range(2) < iatom) cycle
214  ALLOCATE (new_point)
215 
216  ! training-point input, calculate descriptor only on one rank
217  IF (mod(iatom - 1, para_env%num_pe) == para_env%mepos) THEN
218  CALL pao_ml_calc_descriptor(pao, &
219  my_particle_set, &
220  qs_kind_set, &
221  cell, &
222  iatom=iatom, &
223  descriptor=new_point%input)
224  END IF
225 
226  ! copy training-point output on first rank
227  IF (para_env%is_source()) THEN
228  nparams = SIZE(xblocks(iatom)%p, 1)
229  ALLOCATE (new_point%output(nparams))
230  new_point%output(:) = xblocks(iatom)%p(:, 1)
231  END IF
232 
233  ! add to linked list
234  ikind = kindsmap(atom2kind(iatom))
235  training_lists(ikind)%npoints = training_lists(ikind)%npoints + 1
236  new_point%next => training_lists(ikind)%head
237  training_lists(ikind)%head => new_point
238  END DO
239 
240  DEALLOCATE (cell, my_particle_set, hmat, kindsmap, positions, atom2kind)
241 
242  END SUBROUTINE add_to_training_list
243 
244 ! **************************************************************************************************
245 !> \brief Make read-in kinds on to atomic-kinds of this run
246 !> \param pao ...
247 !> \param qs_env ...
248 !> \param kinds ...
249 !> \param kindsmap ...
250 ! **************************************************************************************************
251  SUBROUTINE match_kinds(pao, qs_env, kinds, kindsmap)
252  TYPE(pao_env_type), POINTER :: pao
253  TYPE(qs_environment_type), POINTER :: qs_env
254  TYPE(pao_iokind_type), DIMENSION(:) :: kinds
255  INTEGER, ALLOCATABLE, DIMENSION(:) :: kindsmap
256 
257  CHARACTER(LEN=default_string_length) :: name
258  INTEGER :: ikind, jkind
259  TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
260 
261  CALL get_qs_env(qs_env, atomic_kind_set=atomic_kind_set)
262 
263  cpassert(.NOT. ALLOCATED(kindsmap))
264  ALLOCATE (kindsmap(SIZE(kinds)))
265  kindsmap(:) = -1
266 
267  DO ikind = 1, SIZE(kinds)
268  DO jkind = 1, SIZE(atomic_kind_set)
269  CALL get_atomic_kind(atomic_kind_set(jkind), name=name)
270  ! match kinds via their name
271  IF (trim(kinds(ikind)%name) .EQ. trim(name)) THEN
272  CALL pao_kinds_ensure_equal(pao, qs_env, jkind, kinds(ikind))
273  kindsmap(ikind) = jkind
274  EXIT
275  END IF
276  END DO
277  END DO
278 
279  IF (any(kindsmap < 1)) &
280  cpabort("PAO: Could not match all kinds from training set")
281  END SUBROUTINE match_kinds
282 
283 ! **************************************************************************************************
284 !> \brief Checks that there is at least one training point per pao-enabled kind
285 !> \param qs_env ...
286 !> \param training_lists ...
287 ! **************************************************************************************************
288  SUBROUTINE sanity_check(qs_env, training_lists)
289  TYPE(qs_environment_type), POINTER :: qs_env
290  TYPE(training_list_type), DIMENSION(:), TARGET :: training_lists
291 
292  INTEGER :: ikind, pao_basis_size
293  TYPE(gto_basis_set_type), POINTER :: basis_set
294  TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
295  TYPE(training_list_type), POINTER :: training_list
296 
297  CALL get_qs_env(qs_env, qs_kind_set=qs_kind_set)
298 
299  DO ikind = 1, SIZE(training_lists)
300  training_list => training_lists(ikind)
301  IF (training_list%npoints > 0) cycle ! it's ok
302  CALL get_qs_kind(qs_kind_set(ikind), basis_set=basis_set, pao_basis_size=pao_basis_size)
303  IF (pao_basis_size /= basis_set%nsgf) & ! if this kind has pao enabled...
304  cpabort("Found no training-points for kind: "//trim(training_list%kindname))
305  END DO
306 
307  END SUBROUTINE sanity_check
308 
309 ! **************************************************************************************************
310 !> \brief Turns the linked lists of training points into matrices
311 !> \param training_lists ...
312 !> \param training_matrices ...
313 !> \param para_env ...
314 ! **************************************************************************************************
315  SUBROUTINE training_list2matrix(training_lists, training_matrices, para_env)
316  TYPE(training_list_type), ALLOCATABLE, &
317  DIMENSION(:), TARGET :: training_lists
318  TYPE(training_matrix_type), ALLOCATABLE, &
319  DIMENSION(:), TARGET :: training_matrices
320  TYPE(mp_para_env_type), POINTER :: para_env
321 
322  INTEGER :: i, ikind, inp_size, ninputs, noutputs, &
323  npoints, out_size
324  TYPE(training_list_type), POINTER :: training_list
325  TYPE(training_matrix_type), POINTER :: training_matrix
326  TYPE(training_point_type), POINTER :: cur_point, prev_point
327 
328  cpassert(ALLOCATED(training_lists) .AND. .NOT. ALLOCATED(training_matrices))
329 
330  ALLOCATE (training_matrices(SIZE(training_lists)))
331 
332  DO ikind = 1, SIZE(training_lists)
333  training_list => training_lists(ikind)
334  training_matrix => training_matrices(ikind)
335  training_matrix%kindname = training_list%kindname ! copy kindname
336  npoints = training_list%npoints ! number of points
337  IF (npoints == 0) THEN
338  ALLOCATE (training_matrix%inputs(0, 0))
339  ALLOCATE (training_matrix%outputs(0, 0))
340  cycle
341  END IF
342 
343  ! figure out size of input and output
344  inp_size = 0; out_size = 0
345  IF (ALLOCATED(training_list%head%input)) &
346  inp_size = SIZE(training_list%head%input)
347  IF (ALLOCATED(training_list%head%output)) &
348  out_size = SIZE(training_list%head%output)
349  CALL para_env%sum(inp_size)
350  CALL para_env%sum(out_size)
351 
352  ! allocate matices to hold all training points
353  ALLOCATE (training_matrix%inputs(inp_size, npoints))
354  ALLOCATE (training_matrix%outputs(out_size, npoints))
355  training_matrix%inputs(:, :) = 0.0_dp
356  training_matrix%outputs(:, :) = 0.0_dp
357 
358  ! loop over all training points, consume linked-list in the process
359  ninputs = 0; noutputs = 0
360  cur_point => training_list%head
361  NULLIFY (training_list%head)
362  DO i = 1, npoints
363  IF (ALLOCATED(cur_point%input)) THEN
364  training_matrix%inputs(:, i) = cur_point%input(:)
365  ninputs = ninputs + 1
366  END IF
367  IF (ALLOCATED(cur_point%output)) THEN
368  training_matrix%outputs(:, i) = cur_point%output(:)
369  noutputs = noutputs + 1
370  END IF
371  ! advance to next entry and deallocate the current one
372  prev_point => cur_point
373  cur_point => cur_point%next
374  DEALLOCATE (prev_point)
375  END DO
376  training_list%npoints = 0 ! list is now empty
377 
378  ! sync training_matrix across ranks
379  CALL para_env%sum(training_matrix%inputs)
380  CALL para_env%sum(training_matrix%outputs)
381 
382  ! sanity check
383  CALL para_env%sum(noutputs)
384  CALL para_env%sum(ninputs)
385  cpassert(noutputs == npoints .AND. ninputs == npoints)
386  END DO
387 
388  END SUBROUTINE training_list2matrix
389 
390 ! **************************************************************************************************
391 !> \brief TODO
392 !> \param ml_prior ...
393 !> \param training_matrices ...
394 ! **************************************************************************************************
395  SUBROUTINE pao_ml_substract_prior(ml_prior, training_matrices)
396  INTEGER, INTENT(IN) :: ml_prior
397  TYPE(training_matrix_type), DIMENSION(:), TARGET :: training_matrices
398 
399  INTEGER :: i, ikind, npoints, out_size
400  TYPE(training_matrix_type), POINTER :: training_matrix
401 
402  DO ikind = 1, SIZE(training_matrices)
403  training_matrix => training_matrices(ikind)
404  out_size = SIZE(training_matrix%outputs, 1)
405  npoints = SIZE(training_matrix%outputs, 2)
406  IF (npoints == 0) cycle
407  ALLOCATE (training_matrix%prior(out_size))
408 
409  ! calculate prior
410  SELECT CASE (ml_prior)
411  CASE (pao_ml_prior_zero)
412  training_matrix%prior(:) = 0.0_dp
413  CASE (pao_ml_prior_mean)
414  training_matrix%prior(:) = sum(training_matrix%outputs, 2)/real(npoints, dp)
415  CASE DEFAULT
416  cpabort("PAO: unknown prior")
417  END SELECT
418 
419  ! subtract prior from all training points
420  DO i = 1, npoints
421  training_matrix%outputs(:, i) = training_matrix%outputs(:, i) - training_matrix%prior
422  END DO
423  END DO
424 
425  END SUBROUTINE pao_ml_substract_prior
426 
427 ! **************************************************************************************************
428 !> \brief Print some statistics about the training set and dump it upon request
429 !> \param pao ...
430 !> \param training_matrices ...
431 ! **************************************************************************************************
432  SUBROUTINE pao_ml_print(pao, training_matrices)
433  TYPE(pao_env_type), POINTER :: pao
434  TYPE(training_matrix_type), DIMENSION(:), TARGET :: training_matrices
435 
436  INTEGER :: i, ikind, n, npoints
437  TYPE(training_matrix_type), POINTER :: training_matrix
438 
439  ! dump training data
440  IF (pao%iw_mldata > 0) THEN
441  DO ikind = 1, SIZE(training_matrices)
442  training_matrix => training_matrices(ikind)
443  npoints = SIZE(training_matrix%outputs, 2)
444  DO i = 1, npoints
445  WRITE (pao%iw_mldata, *) "PAO|ML| training-point kind: ", trim(training_matrix%kindname), &
446  " point:", i, " in:", training_matrix%inputs(:, i), &
447  " out:", training_matrix%outputs(:, i)
448  END DO
449  END DO
450  CALL m_flush(pao%iw_mldata)
451  END IF
452 
453  ! print stats
454  IF (pao%iw > 0) THEN
455  DO ikind = 1, SIZE(training_matrices)
456  training_matrix => training_matrices(ikind)
457  n = SIZE(training_matrix%inputs)
458  IF (n == 0) cycle
459  WRITE (pao%iw, "(A,I3,A,E10.1,1X,E10.1,1X,E10.1)") " PAO|ML| Descriptor for kind: "// &
460  trim(training_matrix%kindname)//" size: ", &
461  SIZE(training_matrix%inputs, 1), " min/mean/max: ", &
462  minval(training_matrix%inputs), &
463  sum(training_matrix%inputs)/real(n, dp), &
464  maxval(training_matrix%inputs)
465  END DO
466  END IF
467 
468  END SUBROUTINE pao_ml_print
469 
470 ! **************************************************************************************************
471 !> \brief Calls the actual learning algorthim to traing on the given matrices
472 !> \param pao ...
473 ! **************************************************************************************************
474  SUBROUTINE pao_ml_train(pao)
475  TYPE(pao_env_type), POINTER :: pao
476 
477  CHARACTER(len=*), PARAMETER :: routinen = 'pao_ml_train'
478 
479  INTEGER :: handle
480 
481  CALL timeset(routinen, handle)
482 
483  SELECT CASE (pao%ml_method)
484  CASE (pao_ml_gp)
485  CALL pao_ml_gp_train(pao)
486  CASE (pao_ml_nn)
487  CALL pao_ml_nn_train(pao)
488  CASE (pao_ml_lazy)
489  ! nothing to do
490  CASE DEFAULT
491  cpabort("PAO: unknown machine learning scheme")
492  END SELECT
493 
494  CALL timestop(handle)
495 
496  END SUBROUTINE pao_ml_train
497 
498 ! **************************************************************************************************
499 !> \brief Fills pao%matrix_X based on machine learning predictions
500 !> \param pao ...
501 !> \param qs_env ...
502 ! **************************************************************************************************
503  SUBROUTINE pao_ml_predict(pao, qs_env)
504  TYPE(pao_env_type), POINTER :: pao
505  TYPE(qs_environment_type), POINTER :: qs_env
506 
507  CHARACTER(len=*), PARAMETER :: routinen = 'pao_ml_predict'
508 
509  INTEGER :: acol, arow, handle, iatom, ikind, natoms
510  REAL(dp), ALLOCATABLE, DIMENSION(:) :: descriptor, variances
511  REAL(dp), DIMENSION(:, :), POINTER :: block_x
512  TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
513  TYPE(cell_type), POINTER :: cell
514  TYPE(dbcsr_iterator_type) :: iter
515  TYPE(mp_para_env_type), POINTER :: para_env
516  TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
517  TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
518 
519  CALL timeset(routinen, handle)
520 
521  CALL get_qs_env(qs_env, &
522  para_env=para_env, &
523  cell=cell, &
524  particle_set=particle_set, &
525  atomic_kind_set=atomic_kind_set, &
526  qs_kind_set=qs_kind_set, &
527  natom=natoms)
528 
529  ! fill matrix_X
530  ALLOCATE (variances(natoms))
531  variances(:) = 0.0_dp
532 !$OMP PARALLEL DEFAULT(NONE) SHARED(pao,qs_env,particle_set,qs_kind_set,cell,variances) &
533 !$OMP PRIVATE(iter,arow,acol,iatom,ikind,descriptor,block_X)
534  CALL dbcsr_iterator_start(iter, pao%matrix_X)
535  DO WHILE (dbcsr_iterator_blocks_left(iter))
536  CALL dbcsr_iterator_next_block(iter, arow, acol, block_x)
537  iatom = arow; cpassert(arow == acol)
538  CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
539  IF (SIZE(block_x) == 0) cycle ! pao disabled for iatom
540 
541  ! calculate descriptor
542  CALL pao_ml_calc_descriptor(pao, &
543  particle_set, &
544  qs_kind_set, &
545  cell, &
546  iatom, &
547  descriptor)
548 
549  ! call actual machine learning for prediction
550  CALL pao_ml_predict_low(pao, ikind=ikind, &
551  descriptor=descriptor, &
552  output=block_x(:, 1), &
553  variance=variances(iatom))
554 
555  DEALLOCATE (descriptor)
556 
557  !add prior
558  block_x(:, 1) = block_x(:, 1) + pao%ml_training_matrices(ikind)%prior
559  END DO
560  CALL dbcsr_iterator_stop(iter)
561 !$OMP END PARALLEL
562 
563  ! print variances
564  CALL para_env%sum(variances)
565  IF (pao%iw_mlvar > 0) THEN
566  DO iatom = 1, natoms
567  WRITE (pao%iw_mlvar, *) "PAO|ML| atom:", iatom, " prediction variance:", variances(iatom)
568  END DO
569  CALL m_flush(pao%iw_mlvar)
570  END IF
571 
572  ! one-line summary
573  IF (pao%iw > 0) WRITE (pao%iw, "(A,E20.10,A,T71,I10)") " PAO|ML| max prediction variance:", &
574  maxval(variances), " for atom:", maxloc(variances)
575 
576  IF (maxval(variances) > pao%ml_tolerance) &
577  cpabort("Variance of prediction above ML_TOLERANCE.")
578 
579  DEALLOCATE (variances)
580 
581  CALL timestop(handle)
582 
583  END SUBROUTINE pao_ml_predict
584 
585 ! **************************************************************************************************
586 !> \brief Queries the actual learning algorthim to make a prediction
587 !> \param pao ...
588 !> \param ikind ...
589 !> \param descriptor ...
590 !> \param output ...
591 !> \param variance ...
592 ! **************************************************************************************************
593  SUBROUTINE pao_ml_predict_low(pao, ikind, descriptor, output, variance)
594  TYPE(pao_env_type), POINTER :: pao
595  INTEGER, INTENT(IN) :: ikind
596  REAL(dp), DIMENSION(:), INTENT(IN) :: descriptor
597  REAL(dp), DIMENSION(:), INTENT(OUT) :: output
598  REAL(dp), INTENT(OUT) :: variance
599 
600  SELECT CASE (pao%ml_method)
601  CASE (pao_ml_gp)
602  CALL pao_ml_gp_predict(pao, ikind, descriptor, output, variance)
603  CASE (pao_ml_nn)
604  CALL pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
605  CASE (pao_ml_lazy)
606  output = 0.0_dp ! let's be really lazy and just rely on the prior
607  variance = 0
608  CASE DEFAULT
609  cpabort("PAO: unknown machine learning scheme")
610  END SELECT
611 
612  END SUBROUTINE pao_ml_predict_low
613 
614 ! **************************************************************************************************
615 !> \brief Calculate forces contributed by machine learning
616 !> \param pao ...
617 !> \param qs_env ...
618 !> \param matrix_G ...
619 !> \param forces ...
620 ! **************************************************************************************************
621  SUBROUTINE pao_ml_forces(pao, qs_env, matrix_G, forces)
622  TYPE(pao_env_type), POINTER :: pao
623  TYPE(qs_environment_type), POINTER :: qs_env
624  TYPE(dbcsr_type) :: matrix_g
625  REAL(dp), DIMENSION(:, :), INTENT(INOUT) :: forces
626 
627  CHARACTER(len=*), PARAMETER :: routinen = 'pao_ml_forces'
628 
629  INTEGER :: acol, arow, handle, iatom, ikind
630  REAL(dp), ALLOCATABLE, DIMENSION(:) :: descr_grad, descriptor
631  REAL(dp), DIMENSION(:, :), POINTER :: block_g
632  TYPE(atomic_kind_type), DIMENSION(:), POINTER :: atomic_kind_set
633  TYPE(cell_type), POINTER :: cell
634  TYPE(dbcsr_iterator_type) :: iter
635  TYPE(mp_para_env_type), POINTER :: para_env
636  TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
637  TYPE(qs_kind_type), DIMENSION(:), POINTER :: qs_kind_set
638 
639  CALL timeset(routinen, handle)
640 
641  CALL get_qs_env(qs_env, &
642  para_env=para_env, &
643  cell=cell, &
644  particle_set=particle_set, &
645  atomic_kind_set=atomic_kind_set, &
646  qs_kind_set=qs_kind_set)
647 
648 !$OMP PARALLEL DEFAULT(NONE) SHARED(pao,matrix_G,particle_set,qs_kind_set,cell) &
649 !$OMP REDUCTION(+:forces) &
650 !$OMP PRIVATE(iter,arow,acol,iatom,ikind,block_G,descriptor,descr_grad)
651  CALL dbcsr_iterator_start(iter, matrix_g)
652  DO WHILE (dbcsr_iterator_blocks_left(iter))
653  CALL dbcsr_iterator_next_block(iter, arow, acol, block_g)
654  iatom = arow; cpassert(arow == acol)
655  CALL get_atomic_kind(particle_set(iatom)%atomic_kind, kind_number=ikind)
656  IF (SIZE(block_g) == 0) cycle ! pao disabled for iatom
657 
658  ! calculate descriptor
659  CALL pao_ml_calc_descriptor(pao, &
660  particle_set, &
661  qs_kind_set, &
662  cell, &
663  iatom=iatom, &
664  descriptor=descriptor)
665 
666  ! calcaulte derivate of machine learning prediction
667  CALL pao_ml_gradient_low(pao, ikind=ikind, &
668  descriptor=descriptor, &
669  outer_deriv=block_g(:, 1), &
670  gradient=descr_grad)
671 
672  ! calculate force contributions from descriptor
673  CALL pao_ml_calc_descriptor(pao, &
674  particle_set, &
675  qs_kind_set, &
676  cell, &
677  iatom=iatom, &
678  descr_grad=descr_grad, &
679  forces=forces)
680 
681  DEALLOCATE (descriptor, descr_grad)
682  END DO
683  CALL dbcsr_iterator_stop(iter)
684 !$OMP END PARALLEL
685 
686  CALL timestop(handle)
687 
688  END SUBROUTINE pao_ml_forces
689 
690 ! **************************************************************************************************
691 !> \brief Calculate gradient of machine learning algorithm
692 !> \param pao ...
693 !> \param ikind ...
694 !> \param descriptor ...
695 !> \param outer_deriv ...
696 !> \param gradient ...
697 ! **************************************************************************************************
698  SUBROUTINE pao_ml_gradient_low(pao, ikind, descriptor, outer_deriv, gradient)
699  TYPE(pao_env_type), POINTER :: pao
700  INTEGER, INTENT(IN) :: ikind
701  REAL(dp), DIMENSION(:), INTENT(IN) :: descriptor, outer_deriv
702  REAL(dp), ALLOCATABLE, DIMENSION(:) :: gradient
703 
704  ALLOCATE (gradient(SIZE(descriptor)))
705 
706  SELECT CASE (pao%ml_method)
707  CASE (pao_ml_gp)
708  CALL pao_ml_gp_gradient(pao, ikind, descriptor, outer_deriv, gradient)
709  CASE (pao_ml_nn)
710  CALL pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
711  CASE (pao_ml_lazy)
712  gradient = 0.0_dp
713  CASE DEFAULT
714  cpabort("PAO: unknown machine learning scheme")
715  END SELECT
716 
717  END SUBROUTINE pao_ml_gradient_low
718 
719 END MODULE pao_ml
atomic_kind_types
Define the atomic kind types and their sub types.
Definition: atomic_kind_types.F:23
atomic_kind_types::get_atomic_kind
subroutine, public get_atomic_kind(atomic_kind, fist_potential, element_symbol, name, mass, kind_number, natom, atom_list, rcov, rvdw, z, qeff, apol, cpol, mm_radius, shell, shell_active, damping)
Get attributes of an atomic kind.
Definition: atomic_kind_types.F:138
basis_set_types
Definition: basis_set_types.F:15
cell_methods
Handles all functions related to the CELL.
Definition: cell_methods.F:15
cell_methods::cell_create
subroutine, public cell_create(cell, hmat, periodic, tag)
allocates and initializes a cell
Definition: cell_methods.F:85
cell_types
Handles all functions related to the CELL.
Definition: cell_types.F:15
kinds
Defines the basic variable types.
Definition: kinds.F:23
kinds::dp
integer, parameter, public dp
Definition: kinds.F:34
kinds::default_string_length
integer, parameter, public default_string_length
Definition: kinds.F:57
kinds::default_path_length
integer, parameter, public default_path_length
Definition: kinds.F:58
machine
Machine interface based on Fortran 2003 and POSIX.
Definition: machine.F:17
machine::m_flush
subroutine, public m_flush(lunit)
flushes units if the &GLOBAL flag is set accordingly
Definition: machine.F:106
message_passing
Interface to the message passing library MPI.
Definition: message_passing.F:23
pao_input
Definition: pao_input.F:8
pao_input::id2str
character(len=11) function, public id2str(id)
Helper routine.
Definition: pao_input.F:216
pao_input::pao_ml_gp
integer, parameter, public pao_ml_gp
Definition: pao_input.F:45
pao_input::pao_ml_prior_mean
integer, parameter, public pao_ml_prior_mean
Definition: pao_input.F:45
pao_input::pao_ml_lazy
integer, parameter, public pao_ml_lazy
Definition: pao_input.F:45
pao_input::pao_rotinv_param
integer, parameter, public pao_rotinv_param
Definition: pao_input.F:45
pao_input::pao_ml_prior_zero
integer, parameter, public pao_ml_prior_zero
Definition: pao_input.F:45
pao_input::pao_ml_nn
integer, parameter, public pao_ml_nn
Definition: pao_input.F:45
pao_io
Routines for reading and writing restart files.
Definition: pao_io.F:12
pao_io::pao_kinds_ensure_equal
subroutine, public pao_kinds_ensure_equal(pao, qs_env, ikind, pao_kind)
Ensure that the kind read from the restart is equal to the kind curretly in use.
Definition: pao_io.F:326
pao_io::pao_read_raw
subroutine, public pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks, ml_range)
Reads a restart file into temporary datastructures.
Definition: pao_io.F:183
pao_ml_descriptor
Feature vectors for describing chemical environments in a rotationally invariant fashion.
Definition: pao_ml_descriptor.F:12
pao_ml_descriptor::pao_ml_calc_descriptor
subroutine, public pao_ml_calc_descriptor(pao, particle_set, qs_kind_set, cell, iatom, descriptor, descr_grad, forces)
Calculates a descriptor for chemical environment of given atom.
Definition: pao_ml_descriptor.F:55
pao_ml_gaussprocess
Gaussian Process implementation.
Definition: pao_ml_gaussprocess.F:12
pao_ml_gaussprocess::pao_ml_gp_gradient
subroutine, public pao_ml_gp_gradient(pao, ikind, descriptor, outer_deriv, gradient)
Calculate gradient of Gaussian process.
Definition: pao_ml_gaussprocess.F:130
pao_ml_gaussprocess::pao_ml_gp_train
subroutine, public pao_ml_gp_train(pao)
Builds the covariance matrix.
Definition: pao_ml_gaussprocess.F:33
pao_ml_gaussprocess::pao_ml_gp_predict
subroutine, public pao_ml_gp_predict(pao, ikind, descriptor, output, variance)
Uses covariance matrix to make prediction.
Definition: pao_ml_gaussprocess.F:81
pao_ml_neuralnet
Neural Network implementation.
Definition: pao_ml_neuralnet.F:12
pao_ml_neuralnet::pao_ml_nn_gradient
subroutine, public pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
Calculate gradient of neural network.
Definition: pao_ml_neuralnet.F:68
pao_ml_neuralnet::pao_ml_nn_train
subroutine, public pao_ml_nn_train(pao)
Trains the neural network on given training points.
Definition: pao_ml_neuralnet.F:121
pao_ml_neuralnet::pao_ml_nn_predict
subroutine, public pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
Uses neural network to make a prediction.
Definition: pao_ml_neuralnet.F:44
pao_ml
Main module for PAO Machine Learning.
Definition: pao_ml.F:12
pao_ml::pao_ml_forces
subroutine, public pao_ml_forces(pao, qs_env, matrix_G, forces)
Calculate forces contributed by machine learning.
Definition: pao_ml.F:622
pao_ml::pao_ml_init
subroutine, public pao_ml_init(pao, qs_env)
Initializes the learning machinery.
Definition: pao_ml.F:85
pao_ml::pao_ml_predict
subroutine, public pao_ml_predict(pao, qs_env)
Fills paomatrix_X based on machine learning predictions.
Definition: pao_ml.F:504
pao_types
Types used by the PAO machinery.
Definition: pao_types.F:12
particle_types
Define the data structure for the particle information.
Definition: particle_types.F:19
qs_environment_types
Definition: qs_environment_types.F:14
qs_environment_types::get_qs_env
subroutine, public get_qs_env(qs_env, atomic_kind_set, qs_kind_set, cell, super_cell, cell_ref, use_ref_cell, kpoints, dft_control, mos, sab_orb, sab_all, qmmm, qmmm_periodic, sac_ae, sac_ppl, sac_lri, sap_ppnl, sab_vdw, sab_scp, sap_oce, sab_lrc, sab_se, sab_xtbe, sab_tbe, sab_core, sab_xb, sab_xtb_nonbond, sab_almo, sab_kp, sab_kp_nosym, particle_set, energy, force, matrix_h, matrix_h_im, matrix_ks, matrix_ks_im, matrix_vxc, run_rtp, rtp, matrix_h_kp, matrix_h_im_kp, matrix_ks_kp, matrix_ks_im_kp, matrix_vxc_kp, kinetic_kp, matrix_s_kp, matrix_w_kp, matrix_s_RI_aux_kp, matrix_s, matrix_s_RI_aux, matrix_w, matrix_p_mp2, matrix_p_mp2_admm, rho, rho_xc, pw_env, ewald_env, ewald_pw, active_space, mpools, input, para_env, blacs_env, scf_control, rel_control, kinetic, qs_charges, vppl, rho_core, rho_nlcc, rho_nlcc_g, ks_env, ks_qmmm_env, wf_history, scf_env, local_particles, local_molecules, distribution_2d, dbcsr_dist, molecule_kind_set, molecule_set, subsys, cp_subsys, oce, local_rho_set, rho_atom_set, task_list, task_list_soft, rho0_atom_set, rho0_mpole, rhoz_set, ecoul_1c, rho0_s_rs, rho0_s_gs, do_kpoints, has_unit_metric, requires_mo_derivs, mo_derivs, mo_loc_history, nkind, natom, nelectron_total, nelectron_spin, efield, neighbor_list_id, linres_control, xas_env, virial, cp_ddapc_env, cp_ddapc_ewald, outer_scf_history, outer_scf_ihistory, x_data, et_coupling, dftb_potential, results, se_taper, se_store_int_env, se_nddo_mpole, se_nonbond_env, admm_env, lri_env, lri_density, exstate_env, ec_env, dispersion_env, gcp_env, vee, rho_external, external_vxc, mask, mp2_env, bs_env, kg_env, WannierCentres, atprop, ls_scf_env, do_transport, transport_env, v_hartree_rspace, s_mstruct_changed, rho_changed, potential_changed, forces_up_to_date, mscfg_env, almo_scf_env, gradient_history, variable_history, embed_pot, spin_embed_pot, polar_env, mos_last_converged, rhs)
Get the QUICKSTEP environment.
Definition: qs_environment_types.F:502
qs_kind_types
Define the quickstep kind type and their sub types.
Definition: qs_kind_types.F:23
qs_kind_types::get_qs_kind
subroutine, public get_qs_kind(qs_kind, basis_set, basis_type, ncgf, nsgf, all_potential, tnadd_potential, gth_potential, sgp_potential, upf_potential, se_parameter, dftb_parameter, xtb_parameter, dftb3_param, zeff, elec_conf, mao, lmax_dftb, alpha_core_charge, ccore_charge, core_charge, core_charge_radius, paw_proj_set, paw_atom, hard_radius, hard0_radius, max_rad_local, covalent_radius, vdw_radius, gpw_r3d_rs_type_forced, harmonics, max_iso_not0, max_s_harm, grid_atom, ngrid_ang, ngrid_rad, lmax_rho0, dft_plus_u_atom, l_of_dft_plus_u, n_of_dft_plus_u, u_minus_j, U_of_dft_plus_u, J_of_dft_plus_u, alpha_of_dft_plus_u, beta_of_dft_plus_u, J0_of_dft_plus_u, occupation_of_dft_plus_u, dispersion, bs_occupation, magnetization, no_optimize, addel, laddel, naddel, orbitals, max_scf, eps_scf, smear, u_ramping, u_minus_j_target, eps_u_ramping, init_u_ramping_each_scf, reltmat, ghost, floating, name, element_symbol, pao_basis_size, pao_potentials, pao_descriptors, nelec)
Get attributes of an atomic kind.
Definition: qs_kind_types.F:451