-
Notifications
You must be signed in to change notification settings - Fork 1
/
manybody_allegro.F
516 lines (474 loc) · 23.9 KB
/
manybody_allegro.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
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
!--------------------------------------------------------------------------------------------------!
! 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 !
!--------------------------------------------------------------------------------------------------!
! **************************************************************************************************
!> \par History
!> allegro implementation
!> \author Gabriele Tocci
! **************************************************************************************************
MODULE manybody_allegro
USE atomic_kind_types, ONLY: atomic_kind_type
USE cell_types, ONLY: cell_type
USE fist_neighbor_list_types, ONLY: fist_neighbor_type,&
neighbor_kind_pairs_type
USE fist_nonbond_env_types, ONLY: allegro_data_type,&
fist_nonbond_env_get,&
fist_nonbond_env_set,&
fist_nonbond_env_type,&
pos_type
USE kinds, ONLY: dp,&
int_8,&
sp
USE message_passing, ONLY: mp_para_env_type
USE pair_potential_types, ONLY: allegro_pot_type,&
allegro_type,&
pair_potential_pp_type,&
pair_potential_single_type
USE particle_types, ONLY: particle_type
USE torch_api, ONLY: torch_dict_create,&
torch_dict_get,&
torch_dict_insert,&
torch_dict_release,&
torch_dict_type,&
torch_model_eval,&
torch_model_freeze,&
torch_model_load
USE util, ONLY: sort
#include "./base/base_uses.f90"
IMPLICIT NONE
PRIVATE
PUBLIC :: setup_allegro_arrays, destroy_allegro_arrays, &
allegro_energy_store_force_virial, allegro_add_force_virial
CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'manybody_allegro'
CONTAINS
! **************************************************************************************************
!> \brief ...
!> \param nonbonded ...
!> \param potparm ...
!> \param glob_loc_list ...
!> \param glob_cell_v ...
!> \param glob_loc_list_a ...
!> \param unique_list_a ...
!> \param cell ...
!> \par History
!> Implementation of the allegro potential - [gtocci] 2023
!> \author Gabriele Tocci - University of Zurich
! **************************************************************************************************
SUBROUTINE setup_allegro_arrays(nonbonded, potparm, glob_loc_list, glob_cell_v, glob_loc_list_a, &
unique_list_a, cell)
TYPE(fist_neighbor_type), POINTER :: nonbonded
TYPE(pair_potential_pp_type), POINTER :: potparm
INTEGER, DIMENSION(:, :), POINTER :: glob_loc_list
REAL(KIND=dp), DIMENSION(:, :), POINTER :: glob_cell_v
INTEGER, DIMENSION(:), POINTER :: glob_loc_list_a, unique_list_a
TYPE(cell_type), POINTER :: cell
CHARACTER(LEN=*), PARAMETER :: routineN = 'setup_allegro_arrays'
INTEGER :: handle, i, iend, igrp, ikind, ilist, &
ipair, istart, jkind, nkinds, nlocal, &
npairs, npairs_tot
INTEGER, ALLOCATABLE, DIMENSION(:) :: temp_unique_list_a, work_list, work_list2
INTEGER, DIMENSION(:, :), POINTER :: list
REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: rwork_list
REAL(KIND=dp), DIMENSION(3) :: cell_v, cvi
TYPE(neighbor_kind_pairs_type), POINTER :: neighbor_kind_pair
TYPE(pair_potential_single_type), POINTER :: pot
CPASSERT(.NOT. ASSOCIATED(glob_loc_list))
CPASSERT(.NOT. ASSOCIATED(glob_loc_list_a))
CPASSERT(.NOT. ASSOCIATED(unique_list_a))
CPASSERT(.NOT. ASSOCIATED(glob_cell_v))
CALL timeset(routineN, handle)
npairs_tot = 0
nkinds = SIZE(potparm%pot, 1)
DO ilist = 1, nonbonded%nlists
neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
npairs = neighbor_kind_pair%npairs
IF (npairs == 0) CYCLE
Kind_Group_Loop1: DO igrp = 1, neighbor_kind_pair%ngrp_kind
istart = neighbor_kind_pair%grp_kind_start(igrp)
iend = neighbor_kind_pair%grp_kind_end(igrp)
ikind = neighbor_kind_pair%ij_kind(1, igrp)
jkind = neighbor_kind_pair%ij_kind(2, igrp)
pot => potparm%pot(ikind, jkind)%pot
npairs = iend - istart + 1
IF (pot%no_mb) CYCLE
DO i = 1, SIZE(pot%type)
IF (pot%type(i) == allegro_type) npairs_tot = npairs_tot + npairs
END DO
END DO Kind_Group_Loop1
END DO
ALLOCATE (work_list(npairs_tot))
ALLOCATE (work_list2(npairs_tot))
ALLOCATE (glob_loc_list(2, npairs_tot))
ALLOCATE (glob_cell_v(3, npairs_tot))
! Fill arrays with data
npairs_tot = 0
DO ilist = 1, nonbonded%nlists
neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
npairs = neighbor_kind_pair%npairs
IF (npairs == 0) CYCLE
Kind_Group_Loop2: DO igrp = 1, neighbor_kind_pair%ngrp_kind
istart = neighbor_kind_pair%grp_kind_start(igrp)
iend = neighbor_kind_pair%grp_kind_end(igrp)
ikind = neighbor_kind_pair%ij_kind(1, igrp)
jkind = neighbor_kind_pair%ij_kind(2, igrp)
list => neighbor_kind_pair%list
cvi = neighbor_kind_pair%cell_vector
pot => potparm%pot(ikind, jkind)%pot
npairs = iend - istart + 1
IF (pot%no_mb) CYCLE
cell_v = MATMUL(cell%hmat, cvi)
DO i = 1, SIZE(pot%type)
! ALLEGRO
IF (pot%type(i) == allegro_type) THEN
DO ipair = 1, npairs
glob_loc_list(:, npairs_tot + ipair) = list(:, istart - 1 + ipair)
glob_cell_v(1:3, npairs_tot + ipair) = cell_v(1:3)
END DO
npairs_tot = npairs_tot + npairs
END IF
END DO
END DO Kind_Group_Loop2
END DO
! Order the arrays w.r.t. the first index of glob_loc_list
CALL sort(glob_loc_list(1, :), npairs_tot, work_list)
DO ipair = 1, npairs_tot
work_list2(ipair) = glob_loc_list(2, work_list(ipair))
END DO
glob_loc_list(2, :) = work_list2
DEALLOCATE (work_list2)
ALLOCATE (rwork_list(3, npairs_tot))
DO ipair = 1, npairs_tot
rwork_list(:, ipair) = glob_cell_v(:, work_list(ipair))
END DO
glob_cell_v = rwork_list
DEALLOCATE (rwork_list)
DEALLOCATE (work_list)
ALLOCATE (glob_loc_list_a(npairs_tot))
glob_loc_list_a = glob_loc_list(1, :)
ALLOCATE (temp_unique_list_a(npairs_tot))
nlocal = 1
temp_unique_list_a(1) = glob_loc_list_a(1)
DO ipair = 2, npairs_tot
IF (glob_loc_list_a(ipair - 1) /= glob_loc_list_a(ipair)) THEN
nlocal = nlocal + 1
temp_unique_list_a(nlocal) = glob_loc_list_a(ipair)
END IF
END DO
ALLOCATE (unique_list_a(nlocal))
unique_list_a(:) = temp_unique_list_a(:nlocal)
DEALLOCATE (temp_unique_list_a)
CALL timestop(handle)
END SUBROUTINE setup_allegro_arrays
! **************************************************************************************************
!> \brief ...
!> \param glob_loc_list ...
!> \param glob_cell_v ...
!> \param glob_loc_list_a ...
!> \param unique_list_a ...
!> \par History
!> Implementation of the allegro potential - [gtocci] 2023
!> \author Gabriele Tocci - University of Zurich
! **************************************************************************************************
SUBROUTINE destroy_allegro_arrays(glob_loc_list, glob_cell_v, glob_loc_list_a, unique_list_a)
INTEGER, DIMENSION(:, :), POINTER :: glob_loc_list
REAL(KIND=dp), DIMENSION(:, :), POINTER :: glob_cell_v
INTEGER, DIMENSION(:), POINTER :: glob_loc_list_a, unique_list_a
IF (ASSOCIATED(glob_loc_list)) THEN
DEALLOCATE (glob_loc_list)
END IF
IF (ASSOCIATED(glob_loc_list_a)) THEN
DEALLOCATE (glob_loc_list_a)
END IF
IF (ASSOCIATED(glob_cell_v)) THEN
DEALLOCATE (glob_cell_v)
END IF
IF (ASSOCIATED(unique_list_a)) THEN
DEALLOCATE (unique_list_a)
END IF
END SUBROUTINE destroy_allegro_arrays
! **************************************************************************************************
!> \brief ...
!> \param nonbonded ...
!> \param particle_set ...
!> \param cell ...
!> \param atomic_kind_set ...
!> \param potparm ...
!> \param allegro ...
!> \param glob_loc_list_a ...
!> \param r_last_update_pbc ...
!> \param pot_allegro ...
!> \param fist_nonbond_env ...
!> \param unique_list_a ...
!> \param para_env ...
!> \param use_virial ...
!> \par History
!> Implementation of the allegro potential - [gtocci] 2023
!> Index mapping of atoms from .xyz to Allegro config.yaml file - [mbilichenko] 2024
!> \author Gabriele Tocci - University of Zurich
! **************************************************************************************************
SUBROUTINE allegro_energy_store_force_virial(nonbonded, particle_set, cell, atomic_kind_set, &
potparm, allegro, glob_loc_list_a, r_last_update_pbc, &
pot_allegro, fist_nonbond_env, unique_list_a, para_env, use_virial)
TYPE(fist_neighbor_type), POINTER :: nonbonded
TYPE(particle_type), POINTER :: particle_set(:)
TYPE(cell_type), POINTER :: cell
TYPE(atomic_kind_type), POINTER :: atomic_kind_set(:)
TYPE(pair_potential_pp_type), POINTER :: potparm
TYPE(allegro_pot_type), POINTER :: allegro
INTEGER, DIMENSION(:), POINTER :: glob_loc_list_a
TYPE(pos_type), DIMENSION(:), POINTER :: r_last_update_pbc
REAL(kind=dp) :: pot_allegro
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
INTEGER, DIMENSION(:), POINTER :: unique_list_a
TYPE(mp_para_env_type), POINTER :: para_env
LOGICAL, INTENT(IN) :: use_virial
CHARACTER(LEN=*), PARAMETER :: routineN = 'allegro_energy_store_force_virial'
INTEGER :: atom_a, atom_b, atom_idx, handle, i, iat, iat_use, iend, ifirst, igrp, ikind, &
ilast, ilist, ipair, istart, iunique, jkind, junique, mpair, n_atoms, n_atoms_use, &
nedges, nloc_size, npairs, nunique
INTEGER(kind=int_8), ALLOCATABLE :: atom_types(:), temp_atom_types(:)
INTEGER(kind=int_8), ALLOCATABLE, DIMENSION(:, :) :: edge_index, t_edge_index, temp_edge_index
INTEGER, ALLOCATABLE, DIMENSION(:) :: work_list
INTEGER, DIMENSION(:, :), POINTER :: list, sort_list
LOGICAL, ALLOCATABLE :: use_atom(:)
REAL(kind=dp) :: drij, rab2_max, rij(3)
REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: edge_cell_shifts, lattice, &
new_edge_cell_shifts, pos
REAL(kind=dp), DIMENSION(3) :: cell_v, cvi
REAL(kind=dp), DIMENSION(:, :), POINTER :: atomic_energy, forces, virial
REAL(kind=dp), DIMENSION(:, :, :), POINTER :: virial3d
REAL(kind=sp), ALLOCATABLE, DIMENSION(:, :) :: lattice_sp, new_edge_cell_shifts_sp, &
pos_sp
REAL(kind=sp), DIMENSION(:, :), POINTER :: atomic_energy_sp, forces_sp
TYPE(allegro_data_type), POINTER :: allegro_data
TYPE(neighbor_kind_pairs_type), POINTER :: neighbor_kind_pair
TYPE(pair_potential_single_type), POINTER :: pot
TYPE(torch_dict_type) :: inputs, outputs
CALL timeset(routineN, handle)
NULLIFY (atomic_energy, forces, atomic_energy_sp, forces_sp, virial3d, virial)
n_atoms = SIZE(particle_set)
ALLOCATE (use_atom(n_atoms))
use_atom = .FALSE.
DO ikind = 1, SIZE(atomic_kind_set)
DO jkind = 1, SIZE(atomic_kind_set)
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= allegro_type) CYCLE
DO iat = 1, n_atoms
IF (particle_set(iat)%atomic_kind%kind_number == ikind .OR. &
particle_set(iat)%atomic_kind%kind_number == jkind) use_atom(iat) = .TRUE.
END DO ! iat
END DO ! i
END DO ! jkind
END DO ! ikind
n_atoms_use = COUNT(use_atom)
! get allegro_data to save force, virial info and to load model
CALL fist_nonbond_env_get(fist_nonbond_env, allegro_data=allegro_data)
IF (.NOT. ASSOCIATED(allegro_data)) THEN
ALLOCATE (allegro_data)
CALL fist_nonbond_env_set(fist_nonbond_env, allegro_data=allegro_data)
NULLIFY (allegro_data%use_indices, allegro_data%force)
CALL torch_model_load(allegro_data%model, pot%set(1)%allegro%allegro_file_name)
CALL torch_model_freeze(allegro_data%model)
END IF
IF (ASSOCIATED(allegro_data%force)) THEN
IF (SIZE(allegro_data%force, 2) /= n_atoms_use) THEN
DEALLOCATE (allegro_data%force, allegro_data%use_indices)
END IF
END IF
IF (.NOT. ASSOCIATED(allegro_data%force)) THEN
ALLOCATE (allegro_data%force(3, n_atoms_use))
ALLOCATE (allegro_data%use_indices(n_atoms_use))
END IF
iat_use = 0
DO iat = 1, n_atoms_use
IF (use_atom(iat)) THEN
iat_use = iat_use + 1
allegro_data%use_indices(iat_use) = iat
END IF
END DO
nedges = 0
ALLOCATE (edge_index(2, SIZE(glob_loc_list_a)))
ALLOCATE (edge_cell_shifts(3, SIZE(glob_loc_list_a)))
ALLOCATE (temp_atom_types(SIZE(glob_loc_list_a)))
DO ilist = 1, nonbonded%nlists
neighbor_kind_pair => nonbonded%neighbor_kind_pairs(ilist)
npairs = neighbor_kind_pair%npairs
IF (npairs == 0) CYCLE
Kind_Group_Loop_Allegro: DO igrp = 1, neighbor_kind_pair%ngrp_kind
istart = neighbor_kind_pair%grp_kind_start(igrp)
iend = neighbor_kind_pair%grp_kind_end(igrp)
ikind = neighbor_kind_pair%ij_kind(1, igrp)
jkind = neighbor_kind_pair%ij_kind(2, igrp)
list => neighbor_kind_pair%list
cvi = neighbor_kind_pair%cell_vector
pot => potparm%pot(ikind, jkind)%pot
DO i = 1, SIZE(pot%type)
IF (pot%type(i) /= allegro_type) CYCLE
rab2_max = pot%set(i)%allegro%rcutsq
cell_v = MATMUL(cell%hmat, cvi)
pot => potparm%pot(ikind, jkind)%pot
allegro => pot%set(i)%allegro
npairs = iend - istart + 1
IF (npairs /= 0) THEN
ALLOCATE (sort_list(2, npairs), work_list(npairs))
sort_list = list(:, istart:iend)
! Sort the list of neighbors, this increases the efficiency for single
! potential contributions
CALL sort(sort_list(1, :), npairs, work_list)
DO ipair = 1, npairs
work_list(ipair) = sort_list(2, work_list(ipair))
END DO
sort_list(2, :) = work_list
! find number of unique elements of array index 1
nunique = 1
DO ipair = 1, npairs - 1
IF (sort_list(1, ipair + 1) /= sort_list(1, ipair)) nunique = nunique + 1
END DO
ipair = 1
junique = sort_list(1, ipair)
ifirst = 1
DO iunique = 1, nunique
atom_a = junique
IF (glob_loc_list_a(ifirst) > atom_a) CYCLE
DO mpair = ifirst, SIZE(glob_loc_list_a)
IF (glob_loc_list_a(mpair) == atom_a) EXIT
END DO
ifirst = mpair
DO mpair = ifirst, SIZE(glob_loc_list_a)
IF (glob_loc_list_a(mpair) /= atom_a) EXIT
END DO
ilast = mpair - 1
nloc_size = 0
IF (ifirst /= 0) nloc_size = ilast - ifirst + 1
DO WHILE (ipair <= npairs)
IF (sort_list(1, ipair) /= junique) EXIT
atom_b = sort_list(2, ipair)
rij(:) = r_last_update_pbc(atom_b)%r(:) - r_last_update_pbc(atom_a)%r(:) + cell_v
drij = DOT_PRODUCT(rij, rij)
ipair = ipair + 1
IF (drij <= rab2_max) THEN
nedges = nedges + 1
edge_index(:, nedges) = [atom_a - 1, atom_b - 1]
edge_cell_shifts(:, nedges) = cvi
END IF
END DO
ifirst = ilast + 1
IF (ipair <= npairs) junique = sort_list(1, ipair)
END DO
DEALLOCATE (sort_list, work_list)
END IF
END DO
END DO Kind_Group_Loop_Allegro
END DO
allegro => pot%set(1)%allegro
ALLOCATE (temp_edge_index(2, nedges))
temp_edge_index(:, :) = edge_index(:, :nedges)
ALLOCATE (new_edge_cell_shifts(3, nedges))
new_edge_cell_shifts(:, :) = edge_cell_shifts(:, :nedges)
DEALLOCATE (edge_cell_shifts)
ALLOCATE (t_edge_index(nedges, 2))
t_edge_index(:, :) = TRANSPOSE(temp_edge_index)
DEALLOCATE (temp_edge_index, edge_index)
ALLOCATE (lattice(3, 3), lattice_sp(3, 3))
lattice(:, :) = cell%hmat/pot%set(1)%allegro%unit_cell_val
lattice_sp(:, :) = REAL(lattice, kind=sp)
iat_use = 0
ALLOCATE (pos(3, n_atoms_use), atom_types(n_atoms_use))
DO iat = 1, n_atoms_use
IF (.NOT. use_atom(iat)) CYCLE
iat_use = iat_use + 1
! Find index of the element based on its position in config.yaml file to have correct mapping
DO i = 1, SIZE(allegro%type_names_torch)
IF (particle_set(iat)%atomic_kind%element_symbol == allegro%type_names_torch(i)) THEN
atom_idx = i - 1
END IF
END DO
atom_types(iat_use) = atom_idx
pos(:, iat) = r_last_update_pbc(iat)%r(:)/allegro%unit_coords_val
END DO
CALL torch_dict_create(inputs)
IF (allegro%do_allegro_sp) THEN
ALLOCATE (new_edge_cell_shifts_sp(3, nedges), pos_sp(3, n_atoms_use))
new_edge_cell_shifts_sp(:, :) = REAL(new_edge_cell_shifts(:, :), kind=sp)
pos_sp(:, :) = REAL(pos(:, :), kind=sp)
DEALLOCATE (pos, new_edge_cell_shifts)
CALL torch_dict_insert(inputs, "pos", pos_sp)
CALL torch_dict_insert(inputs, "edge_cell_shift", new_edge_cell_shifts_sp)
CALL torch_dict_insert(inputs, "cell", lattice_sp)
ELSE
CALL torch_dict_insert(inputs, "pos", pos)
CALL torch_dict_insert(inputs, "edge_cell_shift", new_edge_cell_shifts)
CALL torch_dict_insert(inputs, "cell", lattice)
END IF
CALL torch_dict_insert(inputs, "edge_index", t_edge_index)
CALL torch_dict_insert(inputs, "atom_types", atom_types)
CALL torch_dict_create(outputs)
CALL torch_model_eval(allegro_data%model, inputs, outputs)
pot_allegro = 0.0_dp
IF (allegro%do_allegro_sp) THEN
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy_sp)
CALL torch_dict_get(outputs, "forces", forces_sp)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
allegro_data%virial(:, :) = virial(:, :)*allegro%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
allegro_data%force(:, :) = REAL(forces_sp(:, :), kind=dp)*allegro%unit_forces_val
DO iat_use = 1, SIZE(unique_list_a)
i = unique_list_a(iat_use)
pot_allegro = pot_allegro + REAL(atomic_energy_sp(1, i), kind=dp)*allegro%unit_energy_val
END DO
DEALLOCATE (forces_sp, atomic_energy_sp, new_edge_cell_shifts_sp, pos_sp)
ELSE
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
CALL torch_dict_get(outputs, "forces", forces)
IF (use_virial) THEN
ALLOCATE (virial(3, 3))
CALL torch_dict_get(outputs, "virial", virial3d)
virial = RESHAPE(virial3d, (/3, 3/))
allegro_data%virial(:, :) = virial(:, :)*allegro%unit_energy_val
DEALLOCATE (virial, virial3d)
END IF
allegro_data%force(:, :) = forces(:, :)*allegro%unit_forces_val
DO iat_use = 1, SIZE(unique_list_a)
i = unique_list_a(iat_use)
pot_allegro = pot_allegro + atomic_energy(1, i)*allegro%unit_energy_val
END DO
DEALLOCATE (forces, atomic_energy, pos, new_edge_cell_shifts)
END IF
CALL torch_dict_release(inputs)
CALL torch_dict_release(outputs)
DEALLOCATE (t_edge_index, atom_types)
IF (use_virial) allegro_data%virial(:, :) = allegro_data%virial/REAL(para_env%num_pe, dp)
CALL timestop(handle)
END SUBROUTINE allegro_energy_store_force_virial
! **************************************************************************************************
!> \brief ...
!> \param fist_nonbond_env ...
!> \param f_nonbond ...
!> \param pv_nonbond ...
!> \param use_virial ...
! **************************************************************************************************
SUBROUTINE allegro_add_force_virial(fist_nonbond_env, f_nonbond, pv_nonbond, use_virial)
TYPE(fist_nonbond_env_type), POINTER :: fist_nonbond_env
REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT) :: f_nonbond, pv_nonbond
LOGICAL, INTENT(IN) :: use_virial
INTEGER :: iat, iat_use
TYPE(allegro_data_type), POINTER :: allegro_data
CALL fist_nonbond_env_get(fist_nonbond_env, allegro_data=allegro_data)
IF (use_virial) THEN
pv_nonbond = pv_nonbond + allegro_data%virial
END IF
DO iat_use = 1, SIZE(allegro_data%use_indices)
iat = allegro_data%use_indices(iat_use)
CPASSERT(iat >= 1 .AND. iat <= SIZE(f_nonbond, 2))
f_nonbond(1:3, iat) = f_nonbond(1:3, iat) + allegro_data%force(1:3, iat_use)
END DO
END SUBROUTINE allegro_add_force_virial
END MODULE manybody_allegro