(git:a660c7f)
Loading...
Searching...
No Matches
skala_gpw_functional.F
Go to the documentation of this file.
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 Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
10! **************************************************************************************************
11MODULE skala_gpw_functional
12 USE cell_types, ONLY: cell_type,&
13 pbc
14 USE cp_array_utils, ONLY: cp_3d_r_cp_type
15 USE cp_log_handling, ONLY: cp_logger_get_default_io_unit
16 USE input_section_types, ONLY: section_get_rval,&
17 section_vals_get_subs_vals,&
18 section_vals_get_subs_vals2,&
19 section_vals_type,&
20 section_vals_val_get
21 USE kinds, ONLY: default_path_length,&
22 dp,&
23 int_8
24 USE message_passing, ONLY: mp_comm_type
25 USE offload_api, ONLY: offload_set_chosen_device
26 USE particle_types, ONLY: particle_type
27 USE pw_grid_types, ONLY: pw_grid_type
28 USE pw_methods, ONLY: pw_scale,&
29 pw_zero
30 USE pw_pool_types, ONLY: pw_pool_type
31 USE pw_types, ONLY: pw_c1d_gs_type,&
32 pw_r3d_rs_type
33 USE qs_grid_atom, ONLY: grid_atom_type
34 USE skala_gpw_features, ONLY: skala_gpw_atom_partition_hard,&
35 skala_gpw_atom_partition_smooth,&
36 skala_gpw_atom_subchunk_count,&
37 skala_gpw_feature_build,&
38 skala_gpw_feature_build_atom_subchunk,&
39 skala_gpw_feature_release,&
40 skala_gpw_feature_type,&
41 skala_gpw_smooth_partition_derivatives
42 USE skala_torch_api, ONLY: skala_torch_model_get_exc,&
43 skala_torch_model_get_exc_density,&
44 skala_torch_model_load,&
45 skala_torch_model_release,&
46 skala_torch_model_type
47 USE string_utilities, ONLY: uppercase
48 USE torch_api, ONLY: &
49 torch_cuda_device_count, torch_cuda_is_available, torch_dict_create, torch_dict_insert, &
50 torch_dict_release, torch_dict_type, torch_tensor_backward_scalar, torch_tensor_data_ptr, &
51 torch_tensor_from_array, torch_tensor_grad, torch_tensor_release, &
52 torch_tensor_to_device_leaf, torch_tensor_type, torch_use_cuda
53 USE xc_rho_cflags_types, ONLY: xc_rho_cflags_type
54 USE xc_rho_set_types, ONLY: xc_rho_set_create,&
55 xc_rho_set_get,&
56 xc_rho_set_release,&
57 xc_rho_set_type,&
58 xc_rho_set_update
59 USE xc_util, ONLY: xc_pw_divergence,&
60 xc_requires_tmp_g
61#include "./base/base_uses.f90"
62
63 IMPLICIT NONE
64
65 PRIVATE
66
67 CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_gpw_functional'
68 INTEGER, PARAMETER, PRIVATE :: atom_chunk_auto_max_rows = 400000, &
69 atom_chunk_auto_min_rows = 100000, &
70 atom_chunk_auto_row_quantum = 100000, &
71 ncollapsed_grad_per_point = 5, ngrad_per_point = 10
76
77 PUBLIC :: ensure_native_skala_grid_scope, get_gauxc_section, skala_gapw_atom_vxc_of_r, &
78 native_skala_gapw_density_partition, skala_gpw_eval, skala_gpw_exc_density, &
79 xc_section_uses_native_skala_grid, xc_section_uses_gauxc_model
80
81 TYPE(skala_torch_model_type), SAVE :: cached_model
82 CHARACTER(len=default_path_length), SAVE :: cached_model_path = ""
83 LOGICAL, SAVE :: cached_model_loaded = .false.
84 INTEGER, SAVE :: cached_model_cuda_device = -3
85 INTEGER, SAVE :: logged_cuda_device = -3, &
86 logged_cuda_device_count = -1, &
87 logged_cuda_nproc = -1, &
88 logged_cuda_request = -3
89
90CONTAINS
91
92! **************************************************************************************************
93!> \brief Return true if the GAUXC subsection requests the CP2K-native GPW grid path.
94!> \param xc_section ...
95!> \return ...
96! **************************************************************************************************
97 FUNCTION xc_section_uses_native_skala_grid(xc_section) RESULT(uses_native_grid)
98 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
99 LOGICAL :: uses_native_grid
100
101 TYPE(section_vals_type), POINTER :: gauxc_section
102
103 uses_native_grid = .false.
104 gauxc_section => get_gauxc_section(xc_section)
105 IF (ASSOCIATED(gauxc_section)) THEN
106 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=uses_native_grid)
107 END IF
108
109 END FUNCTION xc_section_uses_native_skala_grid
110
111! **************************************************************************************************
112!> \brief Return true if the GAUXC subsection requests a model evaluation.
113!> \param xc_section ...
114!> \return ...
115! **************************************************************************************************
116 FUNCTION xc_section_uses_gauxc_model(xc_section) RESULT(uses_gauxc_model)
117 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
118 LOGICAL :: uses_gauxc_model
119
120 CHARACTER(len=default_path_length) :: model_key, model_name, xc_key, xc_name
121 TYPE(section_vals_type), POINTER :: gauxc_section
122
123 uses_gauxc_model = .false.
124 gauxc_section => get_gauxc_section(xc_section)
125 IF (ASSOCIATED(gauxc_section)) THEN
126 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
127 CALL section_vals_val_get(gauxc_section, "FUNCTIONAL", c_val=xc_name)
128 model_key = adjustl(model_name)
129 xc_key = adjustl(xc_name)
130 CALL uppercase(model_key)
131 CALL uppercase(xc_key)
132 uses_gauxc_model = (trim(model_key) /= "" .AND. trim(model_key) /= "NONE" .AND. &
133 trim(model_key) /= trim(xc_key))
134 END IF
135
136 END FUNCTION xc_section_uses_gauxc_model
137
138! **************************************************************************************************
139!> \brief Return the hard/soft GAPW one-center density partition for native SKALA.
140!> \param xc_section ...
141!> \return ...
142! **************************************************************************************************
143 FUNCTION native_skala_gapw_density_partition(xc_section) RESULT(partition)
144 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
145 INTEGER :: partition
146
147 TYPE(section_vals_type), POINTER :: gauxc_section
148
149 partition = skala_gapw_density_partition_hard_minus_soft
150 gauxc_section => get_gauxc_section(xc_section)
151 IF (ASSOCIATED(gauxc_section)) THEN
152 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_GAPW_DENSITY_PARTITION", &
153 i_val=partition)
154 END IF
155
156 SELECT CASE (partition)
157 CASE (skala_gapw_density_partition_hard_minus_soft, &
158 skala_gapw_density_partition_hard_only, &
159 skala_gapw_density_partition_soft_only, &
160 skala_gapw_density_partition_none)
161 CONTINUE
162 CASE DEFAULT
163 CALL cp_abort(__location__, &
164 "Unknown GAUXC%NATIVE_GRID_GAPW_DENSITY_PARTITION value.")
165 END SELECT
166
167 END FUNCTION native_skala_gapw_density_partition
168
169! **************************************************************************************************
170!> \brief Enforce the currently implemented native SKALA GPW input scope.
171!> \param xc_section ...
172! **************************************************************************************************
173 SUBROUTINE ensure_native_skala_grid_scope(xc_section)
174 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
175
176 CHARACTER(len=default_path_length) :: model_key, model_name
177 INTEGER :: ifun, nfun
178 LOGICAL :: native_grid
179 TYPE(section_vals_type), POINTER :: functionals, gauxc_section, xc_fun
180
181 NULLIFY (gauxc_section)
182 IF (.NOT. ASSOCIATED(xc_section)) THEN
183 cpabort("Native SKALA GPW requires an XC section")
184 END IF
185
186 functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
187 IF (.NOT. ASSOCIATED(functionals)) THEN
188 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL section")
189 END IF
190
191 nfun = 0
192 ifun = 0
193 DO
194 ifun = ifun + 1
195 xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
196 IF (.NOT. ASSOCIATED(xc_fun)) EXIT
197 nfun = nfun + 1
198 IF (xc_fun%section%name == "GAUXC") gauxc_section => xc_fun
199 END DO
200
201 IF (.NOT. ASSOCIATED(gauxc_section)) THEN
202 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
203 END IF
204 IF (nfun /= 1) THEN
205 cpabort("Native SKALA GPW requires GAUXC to be the only XC functional")
206 END IF
207
208 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID", l_val=native_grid)
209 IF (.NOT. native_grid) RETURN
210
211 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_name)
212 model_key = adjustl(model_name)
213 CALL uppercase(model_key)
214 IF (trim(model_key) == "NONE" .OR. trim(model_key) == "") THEN
215 cpabort("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
216 END IF
217
218 END SUBROUTINE ensure_native_skala_grid_scope
219
220! **************************************************************************************************
221!> \brief Evaluate SKALA energy and first derivatives on a CP2K GPW grid.
222!> \param vxc_rho ...
223!> \param vxc_tau ...
224!> \param exc ...
225!> \param rho_r ...
226!> \param rho_g ...
227!> \param tau ...
228!> \param xc_section ...
229!> \param weights ...
230!> \param pw_pool ...
231!> \param particle_set ...
232!> \param cell ...
233!> \param compute_virial ...
234!> \param virial_xc ...
235!> \param just_energy ...
236!> \param atom_force ...
237! **************************************************************************************************
238 SUBROUTINE skala_gpw_eval(vxc_rho, vxc_tau, exc, rho_r, rho_g, tau, xc_section, &
239 weights, pw_pool, particle_set, cell, compute_virial, virial_xc, &
240 just_energy, atom_force)
241 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: vxc_rho, vxc_tau
242 REAL(kind=dp), INTENT(OUT) :: exc
243 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
244 TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER :: rho_g
245 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: tau
246 TYPE(section_vals_type), POINTER :: xc_section
247 TYPE(pw_r3d_rs_type), POINTER :: weights
248 TYPE(pw_pool_type), POINTER :: pw_pool
249 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
250 TYPE(cell_type), POINTER :: cell
251 LOGICAL, INTENT(IN) :: compute_virial
252 REAL(kind=dp), DIMENSION(3, 3), INTENT(OUT) :: virial_xc
253 LOGICAL, INTENT(IN), OPTIONAL :: just_energy
254 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT), &
255 OPTIONAL :: atom_force
256
257 CHARACTER(len=default_path_length) :: model_path
258 INTEGER :: iw, native_grid_atom_chunk_max_rows, native_grid_atom_partition, &
259 native_grid_atom_subchunks, native_grid_cuda_device, nspins, phase_handle, &
260 selected_cuda_device, xc_deriv_method_id, xc_rho_smooth_id
261 LOGICAL :: has_atom_chunk_work, have_atom_coord_grad, lsd, my_just_energy, &
262 native_grid_atom_chunk_routing, native_grid_atom_chunks, native_grid_diagnostics, &
263 native_grid_use_cuda, needs_atom_force, use_atom_subchunks
264 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: density_grad, kin_grad
265 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad_grad
266 REAL(kind=dp), DIMENSION(3, 3) :: virial_before
267 TYPE(section_vals_type), POINTER :: gauxc_section
268 TYPE(skala_gpw_feature_type) :: features
269 TYPE(torch_tensor_type) :: atom_coord_grad_t, &
270 atomic_grid_weight_grad_t, exc_tensor, &
271 grid_coord_grad_t, grid_weight_grad_t
272 TYPE(xc_rho_cflags_type) :: needs
273 TYPE(xc_rho_set_type) :: rho_set
274
275 virial_xc = 0.0_dp
276 exc = 0.0_dp
277 my_just_energy = .false.
278 IF (PRESENT(just_energy)) my_just_energy = just_energy
279 needs_atom_force = PRESENT(atom_force)
280 IF (needs_atom_force) atom_force = 0.0_dp
281 have_atom_coord_grad = .false.
282
283 IF (compute_virial .AND. my_just_energy) THEN
284 CALL cp_abort(__location__, &
285 "Native SKALA GPW stress/virial requires feature gradients.")
286 END IF
287 IF (.NOT. ASSOCIATED(rho_g)) THEN
288 CALL cp_abort(__location__, &
289 "Native SKALA GPW requires the reciprocal-space density to form density gradients.")
290 END IF
291 IF (.NOT. ASSOCIATED(tau)) THEN
292 CALL cp_abort(__location__, &
293 "Native SKALA GPW requires the kinetic-energy density.")
294 END IF
295
296 nspins = SIZE(rho_r)
297 lsd = (nspins /= 1)
298 CALL get_skala_model_path(xc_section, model_path)
299 gauxc_section => get_gauxc_section(xc_section)
300 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
301 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
302 i_val=native_grid_cuda_device)
303 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
304 l_val=native_grid_atom_chunks)
305 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_ROUTING", &
306 l_val=native_grid_atom_chunk_routing)
307 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_MAX_ROWS", &
308 i_val=native_grid_atom_chunk_max_rows)
309 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
310 i_val=native_grid_atom_partition)
311 SELECT CASE (native_grid_atom_partition)
312 CASE (1)
313 native_grid_atom_partition = skala_gpw_atom_partition_hard
314 CASE (2)
315 native_grid_atom_partition = skala_gpw_atom_partition_smooth
316 CASE DEFAULT
317 CALL cp_abort(__location__, &
318 "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
319 END SELECT
320 native_grid_atom_chunk_routing = native_grid_atom_chunk_routing .OR. native_grid_atom_chunks
321 native_grid_atom_chunks = native_grid_atom_chunks .OR. native_grid_atom_chunk_routing
322 IF (native_grid_atom_chunk_max_rows < -1) THEN
323 CALL cp_abort(__location__, &
324 "GAUXC%NATIVE_GRID_ATOM_CHUNK_MAX_ROWS must be -1, zero, or positive.")
325 END IF
326 IF (needs_atom_force .OR. compute_virial) THEN
327 IF (native_grid_atom_partition == skala_gpw_atom_partition_hard) THEN
328 native_grid_atom_partition = skala_gpw_atom_partition_smooth
329 END IF
330 native_grid_atom_chunk_routing = .false.
331 native_grid_atom_chunks = .false.
332 END IF
333 ! The portable SKALA export used by the regtests builds ragged-index tensors on CPU.
334 CALL torch_use_cuda(native_grid_use_cuda)
335 selected_cuda_device = configure_native_grid_cuda( &
336 native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
337 CALL ensure_model_loaded(model_path, selected_cuda_device)
338
339 IF (lsd) THEN
340 needs%rho_spin = .true.
341 needs%drho_spin = .true.
342 needs%tau_spin = .true.
343 ELSE
344 needs%rho = .true.
345 needs%drho = .true.
346 needs%tau = .true.
347 END IF
348
349 CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
350 CALL section_vals_val_get(xc_section, "XC_GRID%XC_SMOOTH_RHO", i_val=xc_rho_smooth_id)
351
352 CALL xc_rho_set_create(rho_set, &
353 rho_r(1)%pw_grid%bounds_local, &
354 rho_cutoff=section_get_rval(xc_section, "density_cutoff"), &
355 drho_cutoff=section_get_rval(xc_section, "gradient_cutoff"), &
356 tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
357 CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
358 xc_deriv_method_id, xc_rho_smooth_id, pw_pool)
359
360 CALL skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
361 requires_grad=(.NOT. my_just_energy), weights=weights, &
362 requires_coordinate_grad=(needs_atom_force .OR. compute_virial), &
363 requires_stress_grad=compute_virial, &
364 use_atom_chunks=native_grid_atom_chunks, &
365 route_atom_chunks=native_grid_atom_chunk_routing, &
366 atom_partition=native_grid_atom_partition)
367 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_DIAGNOSTICS", l_val=native_grid_diagnostics)
368 IF (native_grid_diagnostics) THEN
369 CALL print_native_grid_diagnostics(features, rho_r(1)%pw_grid%para%group%mepos == 0)
370 END IF
371
372 IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows == -1) THEN
373 IF (native_grid_use_cuda) THEN
374 native_grid_atom_chunk_max_rows = auto_atom_chunk_max_rows(features, &
375 rho_r(1)%pw_grid%para%group)
376 ELSE
377 native_grid_atom_chunk_max_rows = 0
378 END IF
379 END IF
380 IF (native_grid_diagnostics .AND. features%uses_atom_chunks .AND. &
381 rho_r(1)%pw_grid%para%group%mepos == 0) THEN
382 iw = cp_logger_get_default_io_unit()
383 IF (iw > 0) THEN
384 WRITE (unit=iw, fmt="(T2,A,1X,I0)") &
385 "SKALA_GPW| Native grid atom chunk max rows", native_grid_atom_chunk_max_rows
386 END IF
387 END IF
388 native_grid_atom_subchunks = 1
389 IF (features%uses_atom_chunks .AND. native_grid_atom_chunk_max_rows > 0) THEN
390 native_grid_atom_subchunks = skala_gpw_atom_subchunk_count(native_grid_atom_chunk_max_rows)
391 CALL rho_r(1)%pw_grid%para%group%max(native_grid_atom_subchunks)
392 END IF
393 use_atom_subchunks = features%uses_atom_chunks .AND. native_grid_atom_subchunks > 1
394 has_atom_chunk_work = .NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0
395 exc = 0.0_dp
396 IF (use_atom_subchunks) THEN
397 CALL evaluate_atom_subchunks(features, rho_r(1)%pw_grid%para%group, &
398 native_grid_atom_chunk_max_rows, &
399 compute_grads=(.NOT. my_just_energy), exc=exc, &
400 density_grad=density_grad, grad_grad=grad_grad, &
401 kin_grad=kin_grad, collapse_spin_grads=(nspins == 1))
402 ELSE IF (has_atom_chunk_work) THEN
403 CALL skala_torch_model_get_exc(cached_model, features%inputs, &
404 features%grid_weights_t, exc_tensor, exc)
405 END IF
406 IF (features%uses_atom_chunks) CALL rho_r(1)%pw_grid%para%group%sum(exc)
407
408 IF (.NOT. my_just_energy) THEN
409 IF (.NOT. use_atom_subchunks) THEN
410 IF (has_atom_chunk_work) THEN
411 CALL timeset("skala_gpw_backward", phase_handle)
412 CALL torch_tensor_backward_scalar(exc_tensor)
413 CALL timestop(phase_handle)
414
415 IF (compute_virial) THEN
416 IF (native_grid_diagnostics) virial_before = virial_xc
417 CALL build_weight_virial(virial_xc, features, exc, grid_weight_grad_t, &
418 atomic_grid_weight_grad_t, &
419 rho_r(1)%pw_grid%para%group%mepos == 0, &
420 native_grid_diagnostics)
421 IF (native_grid_diagnostics) THEN
422 CALL print_virial_delta("weight-residual", virial_xc - virial_before, &
423 rho_r(1)%pw_grid%para%group%mepos == 0)
424 END IF
425 END IF
426 END IF
427
428 CALL timeset("skala_gpw_grad_fetch", phase_handle)
429 IF (features%uses_atom_chunks) THEN
430 CALL fetch_and_gather_atom_chunk_grads(features, rho_r(1)%pw_grid%para%group, &
431 density_grad, grad_grad, kin_grad)
432 ELSE
433 CALL fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
434 END IF
435 CALL timestop(phase_handle)
436 END IF
437 IF (needs_atom_force) THEN
438 CALL add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, &
439 rho_r(1)%pw_grid%para%group%mepos == 0)
440 IF (features%atom_partition == skala_gpw_atom_partition_smooth) THEN
441 CALL add_smooth_partition_force(atom_force, features, particle_set, cell, rho_r, &
442 grid_weight_grad_t, atomic_grid_weight_grad_t)
443 END IF
444 have_atom_coord_grad = .true.
445 END IF
446
447 CALL timeset("skala_gpw_vxc_unpack", phase_handle)
448 IF (compute_virial) THEN
449 IF (native_grid_diagnostics) virial_before = virial_xc
450 CALL build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
451 IF (native_grid_diagnostics) THEN
452 CALL print_virial_delta("feature-gradient", virial_xc - virial_before, &
453 rho_r(1)%pw_grid%para%group%mepos == 0)
454 virial_before = virial_xc
455 END IF
456 IF (.NOT. have_atom_coord_grad) THEN
457 CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
458 have_atom_coord_grad = .true.
459 END IF
460 CALL build_static_coordinate_virial(virial_xc, features, atom_coord_grad_t, &
461 grid_coord_grad_t, &
462 rho_r(1)%pw_grid%para%group%mepos == 0, &
463 native_grid_diagnostics)
464 IF (native_grid_diagnostics) THEN
465 CALL print_virial_delta("static-coordinates", virial_xc - virial_before, &
466 rho_r(1)%pw_grid%para%group%mepos == 0)
467 virial_before = virial_xc
468 END IF
469 IF (features%atom_partition == skala_gpw_atom_partition_smooth) THEN
470 CALL build_smooth_partition_virial(virial_xc, features, particle_set, cell, rho_r, &
471 grid_weight_grad_t, atomic_grid_weight_grad_t)
472 IF (native_grid_diagnostics) THEN
473 CALL print_virial_delta("smooth-partition", virial_xc - virial_before, &
474 rho_r(1)%pw_grid%para%group%mepos == 0)
475 virial_before = virial_xc
476 END IF
477 END IF
478 END IF
479 CALL build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
480 density_grad, grad_grad, kin_grad, &
481 xc_deriv_method_id)
482 CALL timestop(phase_handle)
483
484 CALL timeset("skala_gpw_grad_release", phase_handle)
485 DEALLOCATE (density_grad, grad_grad, kin_grad)
486 IF (have_atom_coord_grad) CALL torch_tensor_release(atom_coord_grad_t)
487 CALL timestop(phase_handle)
488 END IF
489
490 CALL timeset("skala_gpw_cleanup", phase_handle)
491 IF (.NOT. use_atom_subchunks .AND. has_atom_chunk_work) CALL torch_tensor_release(exc_tensor)
492 CALL skala_gpw_feature_release(features)
493 CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
494 CALL torch_use_cuda(.true.)
495 CALL timestop(phase_handle)
496
497 END SUBROUTINE skala_gpw_eval
498
499! **************************************************************************************************
500!> \brief Evaluate the native SKALA XC energy density on the CP2K PW grid.
501!> \param exc_r ...
502!> \param rho_r ...
503!> \param rho_g ...
504!> \param tau ...
505!> \param xc_section ...
506!> \param weights ...
507!> \param pw_pool ...
508!> \param particle_set ...
509!> \param cell ...
510! **************************************************************************************************
511 SUBROUTINE skala_gpw_exc_density(exc_r, rho_r, rho_g, tau, xc_section, weights, pw_pool, &
512 particle_set, cell)
513 TYPE(pw_r3d_rs_type), INTENT(INOUT) :: exc_r
514 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
515 TYPE(pw_c1d_gs_type), DIMENSION(:), POINTER :: rho_g
516 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: tau
517 TYPE(section_vals_type), POINTER :: xc_section
518 TYPE(pw_r3d_rs_type), POINTER :: weights
519 TYPE(pw_pool_type), POINTER :: pw_pool
520 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
521 TYPE(cell_type), POINTER :: cell
522
523 CHARACTER(len=default_path_length) :: model_path
524 INTEGER :: feature_pos, i, j, k, local_row, native_grid_atom_partition, &
525 native_grid_cuda_device, nspins, row, selected_cuda_device, xc_deriv_method_id, &
526 xc_rho_smooth_id
527 LOGICAL :: lsd, native_grid_atom_chunk_routing, &
528 native_grid_atom_chunks, &
529 native_grid_use_cuda
530 REAL(kind=dp) :: local_exc
531 REAL(kind=dp), DIMENSION(:), POINTER :: exc_density
532 TYPE(section_vals_type), POINTER :: gauxc_section
533 TYPE(skala_gpw_feature_type) :: features
534 TYPE(torch_tensor_type) :: exc_density_t
535 TYPE(xc_rho_cflags_type) :: needs
536 TYPE(xc_rho_set_type) :: rho_set
537
538 cpassert(ASSOCIATED(rho_r))
539 cpassert(ASSOCIATED(rho_g))
540 cpassert(ASSOCIATED(tau))
541 CALL pw_zero(exc_r)
542
543 nspins = SIZE(rho_r)
544 lsd = (nspins /= 1)
545 CALL get_skala_model_path(xc_section, model_path)
546 gauxc_section => get_gauxc_section(xc_section)
547 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
548 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
549 i_val=native_grid_cuda_device)
550 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNKS", &
551 l_val=native_grid_atom_chunks)
552 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_CHUNK_ROUTING", &
553 l_val=native_grid_atom_chunk_routing)
554 native_grid_atom_chunks = .false.
555 native_grid_atom_chunk_routing = .false.
556 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_ATOM_PARTITION", &
557 i_val=native_grid_atom_partition)
558 SELECT CASE (native_grid_atom_partition)
559 CASE (1)
560 native_grid_atom_partition = skala_gpw_atom_partition_hard
561 CASE (2)
562 native_grid_atom_partition = skala_gpw_atom_partition_smooth
563 CASE DEFAULT
564 CALL cp_abort(__location__, &
565 "Unknown GAUXC%NATIVE_GRID_ATOM_PARTITION value.")
566 END SELECT
567
568 CALL torch_use_cuda(native_grid_use_cuda)
569 selected_cuda_device = configure_native_grid_cuda( &
570 native_grid_use_cuda, native_grid_cuda_device, rho_r(1)%pw_grid%para%group)
571 CALL ensure_model_loaded(model_path, selected_cuda_device)
572
573 IF (lsd) THEN
574 needs%rho_spin = .true.
575 needs%drho_spin = .true.
576 needs%tau_spin = .true.
577 ELSE
578 needs%rho = .true.
579 needs%drho = .true.
580 needs%tau = .true.
581 END IF
582
583 CALL section_vals_val_get(xc_section, "XC_GRID%XC_DERIV", i_val=xc_deriv_method_id)
584 CALL section_vals_val_get(xc_section, "XC_GRID%XC_SMOOTH_RHO", i_val=xc_rho_smooth_id)
585
586 CALL xc_rho_set_create(rho_set, &
587 rho_r(1)%pw_grid%bounds_local, &
588 rho_cutoff=section_get_rval(xc_section, "density_cutoff"), &
589 drho_cutoff=section_get_rval(xc_section, "gradient_cutoff"), &
590 tau_cutoff=section_get_rval(xc_section, "tau_cutoff"))
591 CALL xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, &
592 xc_deriv_method_id, xc_rho_smooth_id, pw_pool)
593
594 CALL skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
595 requires_grad=.false., weights=weights, &
596 requires_coordinate_grad=.false., &
597 requires_stress_grad=.false., &
598 use_atom_chunks=.false., route_atom_chunks=.false., &
599 atom_partition=native_grid_atom_partition)
600 CALL skala_torch_model_get_exc_density(cached_model, features%inputs, exc_density_t)
601 NULLIFY (exc_density)
602 CALL torch_tensor_data_ptr(exc_density_t, exc_density)
603
604 local_row = 0
605 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
606 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
607 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
608 local_row = local_row + 1
609 local_exc = 0.0_dp
610 DO feature_pos = features%local_feature_offsets(local_row), &
611 features%local_feature_offsets(local_row + 1) - 1
612 row = features%local_feature_rows(feature_pos)
613 local_exc = local_exc + exc_density(row)*features%grid_weights(row)
614 END DO
615 exc_r%array(i, j, k) = local_exc/rho_r(1)%pw_grid%dvol
616 END DO
617 END DO
618 END DO
619 cpassert(local_row == features%nflat_local)
620
621 CALL torch_tensor_release(exc_density_t)
622 CALL skala_gpw_feature_release(features)
623 CALL xc_rho_set_release(rho_set, pw_pool=pw_pool)
624 CALL torch_use_cuda(.true.)
625
626 END SUBROUTINE skala_gpw_exc_density
627
628! **************************************************************************************************
629!> \brief Evaluate SKALA on a GAPW one-center atomic grid.
630!> \param xc_section ...
631!> \param grid_atom ...
632!> \param group ...
633!> \param atom_coord ...
634!> \param rho ...
635!> \param drho ...
636!> \param tau ...
637!> \param weights ...
638!> \param lsd ...
639!> \param nspins ...
640!> \param na ...
641!> \param nr ...
642!> \param exc ...
643!> \param vxc ...
644!> \param vxg ...
645!> \param vtau ...
646!> \param energy_only ...
647!> \param atom_force ...
648!> \param atom_virial ...
649! **************************************************************************************************
650 SUBROUTINE skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, &
651 rho, drho, tau, weights, lsd, nspins, na, nr, &
652 exc, vxc, vxg, vtau, energy_only, atom_force, atom_virial)
653 TYPE(section_vals_type), POINTER :: xc_section
654 TYPE(grid_atom_type), POINTER :: grid_atom
655
656 CLASS(mp_comm_type), INTENT(IN) :: group
657 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: atom_coord
658 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: rho, tau, vxc, vtau
659 REAL(kind=dp), DIMENSION(:, :, :, :), POINTER :: drho, vxg
660 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: weights
661 LOGICAL, INTENT(IN) :: lsd
662 INTEGER, INTENT(IN) :: nspins, na, nr
663 REAL(kind=dp), INTENT(OUT) :: exc
664 LOGICAL, INTENT(IN), OPTIONAL :: energy_only
665 REAL(kind=dp), DIMENSION(3), INTENT(OUT), &
666 OPTIONAL :: atom_force
667 REAL(kind=dp), DIMENSION(3, 3), INTENT(OUT), &
668 OPTIONAL :: atom_virial
669
670 CHARACTER(len=default_path_length) :: model_path
671 INTEGER :: ia, idir, ir, native_grid_cuda_device, &
672 jdir, nflat, row, selected_cuda_device
673 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
674 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
675 LOGICAL :: need_coord_grad, my_energy_only, native_grid_use_cuda
676 REAL(kind=dp) :: tmp
677 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
678 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: coarse_0_atomic_coords, density, &
679 grid_coords, kin
680 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: grad
681 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, density_grad, &
682 grid_coord_grad, kin_grad
683 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: grad_grad
684 TYPE(section_vals_type), POINTER :: gauxc_section
685 TYPE(torch_dict_type) :: inputs
686 TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t, &
687 atomic_grid_sizes_t, &
688 atomic_grid_weights_t, &
689 atom_coord_grad_t, &
690 coarse_0_atomic_coords_t, density_t, &
691 density_grad_t, exc_tensor, grad_t, &
692 grad_grad_t, grid_coord_grad_t, &
693 grid_coords_t, grid_weights_t, kin_t, &
694 kin_grad_t
695
696 cpassert(ASSOCIATED(xc_section))
697 cpassert(ASSOCIATED(grid_atom))
698 cpassert(ASSOCIATED(rho))
699 cpassert(ASSOCIATED(drho))
700 cpassert(ASSOCIATED(tau))
701
702 my_energy_only = .false.
703 IF (PRESENT(energy_only)) my_energy_only = energy_only
704 need_coord_grad = PRESENT(atom_force) .OR. PRESENT(atom_virial)
705 exc = 0.0_dp
706 IF (PRESENT(atom_force)) atom_force = 0.0_dp
707 IF (PRESENT(atom_virial)) atom_virial = 0.0_dp
708 IF (.NOT. my_energy_only) THEN
709 vxc = 0.0_dp
710 vxg = 0.0_dp
711 vtau = 0.0_dp
712 END IF
713
714 CALL get_skala_model_path(xc_section, model_path)
715 gauxc_section => get_gauxc_section(xc_section)
716 cpassert(ASSOCIATED(gauxc_section))
717 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
718 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_CUDA_DEVICE", &
719 i_val=native_grid_cuda_device)
720 CALL torch_use_cuda(native_grid_use_cuda)
721 selected_cuda_device = configure_native_grid_cuda( &
722 native_grid_use_cuda, native_grid_cuda_device, group)
723 CALL ensure_model_loaded(model_path, selected_cuda_device)
724
725 nflat = na*nr
726 ALLOCATE (density(nflat, 2), grad(nflat, 3, 2), kin(nflat, 2), &
727 grid_coords(3, nflat), grid_weights(nflat), &
728 atomic_grid_weights(nflat), atomic_grid_sizes(1), &
729 coarse_0_atomic_coords(3, 1), atomic_grid_size_bound_shape(0, nflat))
730 density = 0.0_dp
731 grad = 0.0_dp
732 kin = 0.0_dp
733 grid_coords = 0.0_dp
734 grid_weights = 0.0_dp
735 atomic_grid_weights = 0.0_dp
736 atomic_grid_sizes(1) = int(nflat, kind=int_8)
737 atomic_grid_size_bound_shape = 0_int_8
738 coarse_0_atomic_coords(:, 1) = atom_coord
739
740 row = 0
741 DO ir = 1, nr
742 DO ia = 1, na
743 row = row + 1
744 grid_coords(1, row) = atom_coord(1) + grid_atom%rad(ir)* &
745 grid_atom%sin_pol(ia)*grid_atom%cos_azi(ia)
746 grid_coords(2, row) = atom_coord(2) + grid_atom%rad(ir)* &
747 grid_atom%sin_pol(ia)*grid_atom%sin_azi(ia)
748 grid_coords(3, row) = atom_coord(3) + grid_atom%rad(ir)*grid_atom%cos_pol(ia)
749 grid_weights(row) = weights(ia, ir)
750 atomic_grid_weights(row) = weights(ia, ir)
751 IF (nspins == 1) THEN
752 density(row, :) = 0.5_dp*rho(ia, ir, 1)
753 DO idir = 1, 3
754 grad(row, idir, :) = 0.5_dp*drho(idir, ia, ir, 1)
755 END DO
756 kin(row, :) = 0.5_dp*tau(ia, ir, 1)
757 ELSE
758 density(row, :) = rho(ia, ir, 1:2)
759 DO idir = 1, 3
760 grad(row, idir, :) = drho(idir, ia, ir, 1:2)
761 END DO
762 kin(row, :) = tau(ia, ir, 1:2)
763 END IF
764 END DO
765 END DO
766
767 CALL torch_tensor_from_array(grid_coords_t, grid_coords)
768 CALL torch_tensor_to_device_leaf(grid_coords_t, need_coord_grad)
769 CALL torch_tensor_from_array(grid_weights_t, grid_weights)
770 CALL torch_tensor_to_device_leaf(grid_weights_t, .false.)
771 CALL torch_tensor_from_array(atomic_grid_weights_t, atomic_grid_weights)
772 CALL torch_tensor_to_device_leaf(atomic_grid_weights_t, .false.)
773 CALL torch_tensor_from_array(atomic_grid_sizes_t, atomic_grid_sizes)
774 CALL torch_tensor_to_device_leaf(atomic_grid_sizes_t, .false.)
775 CALL torch_tensor_from_array(atomic_grid_size_bound_shape_t, &
776 atomic_grid_size_bound_shape)
777 CALL torch_tensor_to_device_leaf(atomic_grid_size_bound_shape_t, .false.)
778 CALL torch_tensor_from_array(coarse_0_atomic_coords_t, coarse_0_atomic_coords)
779 CALL torch_tensor_to_device_leaf(coarse_0_atomic_coords_t, need_coord_grad)
780 CALL torch_tensor_from_array(density_t, density)
781 CALL torch_tensor_to_device_leaf(density_t,.NOT. my_energy_only)
782 CALL torch_tensor_from_array(grad_t, grad)
783 CALL torch_tensor_to_device_leaf(grad_t,.NOT. my_energy_only)
784 CALL torch_tensor_from_array(kin_t, kin)
785 CALL torch_tensor_to_device_leaf(kin_t,.NOT. my_energy_only)
786
787 CALL torch_dict_create(inputs)
788 CALL torch_dict_insert(inputs, "grid_coords", grid_coords_t)
789 CALL torch_dict_insert(inputs, "grid_weights", grid_weights_t)
790 CALL torch_dict_insert(inputs, "atomic_grid_weights", atomic_grid_weights_t)
791 CALL torch_dict_insert(inputs, "atomic_grid_sizes", atomic_grid_sizes_t)
792 CALL torch_dict_insert(inputs, "atomic_grid_size_bound_shape", &
793 atomic_grid_size_bound_shape_t)
794 CALL torch_dict_insert(inputs, "density", density_t)
795 CALL torch_dict_insert(inputs, "grad", grad_t)
796 CALL torch_dict_insert(inputs, "kin", kin_t)
797 CALL torch_dict_insert(inputs, "coarse_0_atomic_coords", coarse_0_atomic_coords_t)
798
799 CALL skala_torch_model_get_exc(cached_model, inputs, grid_weights_t, exc_tensor, exc)
800
801 IF (.NOT. my_energy_only) THEN
802 NULLIFY (atom_coord_grad, density_grad, grad_grad, grid_coord_grad, kin_grad)
803 CALL torch_tensor_backward_scalar(exc_tensor)
804 IF (need_coord_grad) THEN
805 CALL torch_tensor_grad(grid_coords_t, grid_coord_grad_t)
806 CALL torch_tensor_grad(coarse_0_atomic_coords_t, atom_coord_grad_t)
807 CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
808 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
809 IF (PRESENT(atom_force)) THEN
810 atom_force(:) = atom_coord_grad(:, 1)
811 DO row = 1, nflat
812 atom_force(:) = atom_force(:) + grid_coord_grad(:, row)
813 END DO
814 END IF
815 IF (PRESENT(atom_virial)) THEN
816 DO row = 1, nflat
817 DO idir = 1, 3
818 DO jdir = 1, 3
819 tmp = grid_coord_grad(idir, row)*coarse_0_atomic_coords(jdir, 1)
820 atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
821 END DO
822 END DO
823 END DO
824 DO idir = 1, 3
825 DO jdir = 1, 3
826 tmp = atom_coord_grad(idir, 1)*coarse_0_atomic_coords(jdir, 1)
827 atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
828 END DO
829 END DO
830 END IF
831 END IF
832 CALL torch_tensor_grad(density_t, density_grad_t)
833 CALL torch_tensor_grad(grad_t, grad_grad_t)
834 CALL torch_tensor_grad(kin_t, kin_grad_t)
835 CALL torch_tensor_data_ptr(density_grad_t, density_grad)
836 CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
837 CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
838
839 row = 0
840 DO ir = 1, nr
841 DO ia = 1, na
842 row = row + 1
843 IF (lsd) THEN
844 vxc(ia, ir, 1:2) = density_grad(row, 1:2)
845 DO idir = 1, 3
846 vxg(idir, ia, ir, 1:2) = grad_grad(row, idir, 1:2)
847 END DO
848 vtau(ia, ir, 1:2) = kin_grad(row, 1:2)
849 ELSE
850 vxc(ia, ir, 1) = 0.5_dp*(density_grad(row, 1) + density_grad(row, 2))
851 DO idir = 1, 3
852 vxg(idir, ia, ir, 1) = &
853 0.5_dp*(grad_grad(row, idir, 1) + grad_grad(row, idir, 2))
854 END DO
855 vtau(ia, ir, 1) = 0.5_dp*(kin_grad(row, 1) + kin_grad(row, 2))
856 END IF
857 END DO
858 END DO
859
860 CALL torch_tensor_release(density_grad_t)
861 CALL torch_tensor_release(grad_grad_t)
862 CALL torch_tensor_release(kin_grad_t)
863 IF (need_coord_grad) THEN
864 CALL torch_tensor_release(grid_coord_grad_t)
865 CALL torch_tensor_release(atom_coord_grad_t)
866 END IF
867 END IF
868
869 CALL torch_tensor_release(exc_tensor)
870 CALL torch_tensor_release(density_t)
871 CALL torch_tensor_release(grad_t)
872 CALL torch_tensor_release(kin_t)
873 CALL torch_tensor_release(grid_coords_t)
874 CALL torch_tensor_release(grid_weights_t)
875 CALL torch_tensor_release(atomic_grid_weights_t)
876 CALL torch_tensor_release(atomic_grid_sizes_t)
877 CALL torch_tensor_release(atomic_grid_size_bound_shape_t)
878 CALL torch_tensor_release(coarse_0_atomic_coords_t)
879 CALL torch_dict_release(inputs)
880 DEALLOCATE (atomic_grid_size_bound_shape, atomic_grid_sizes, atomic_grid_weights, &
881 coarse_0_atomic_coords, density, grad, grid_coords, grid_weights, kin)
882 CALL torch_use_cuda(.true.)
883
884 END SUBROUTINE skala_gapw_atom_vxc_of_r
885
886! **************************************************************************************************
887!> \brief Add the explicit SKALA derivative with respect to atom-center coordinates.
888!> \param atom_force ...
889!> \param features ...
890!> \param atom_coord_grad_t ...
891!> \param root_rank ...
892! **************************************************************************************************
893 SUBROUTINE add_explicit_coordinate_force(atom_force, features, atom_coord_grad_t, root_rank)
894 REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT) :: atom_force
895 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
896 TYPE(torch_tensor_type), INTENT(INOUT) :: atom_coord_grad_t
897 LOGICAL, INTENT(IN) :: root_rank
898
899 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad
900
901 NULLIFY (atom_coord_grad)
902 CALL torch_tensor_grad(features%coarse_0_atomic_coords_t, atom_coord_grad_t)
903 IF (root_rank) THEN
904 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
905 cpassert(SIZE(atom_force, 1) == SIZE(atom_coord_grad, 1))
906 cpassert(SIZE(atom_force, 2) == SIZE(atom_coord_grad, 2))
907 atom_force(:, :) = atom_force(:, :) + atom_coord_grad(:, :)
908 END IF
909
910 END SUBROUTINE add_explicit_coordinate_force
911
912! **************************************************************************************************
913!> \brief Add the force from SMOOTH native-grid atom partition weights.
914!> \param atom_force ...
915!> \param features ...
916!> \param particle_set ...
917!> \param cell ...
918!> \param rho_r ...
919!> \param grid_weight_grad_t ...
920!> \param atomic_grid_weight_grad_t ...
921! **************************************************************************************************
922 SUBROUTINE add_smooth_partition_force(atom_force, features, particle_set, cell, rho_r, &
923 grid_weight_grad_t, atomic_grid_weight_grad_t)
924 REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT) :: atom_force
925 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
926 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
927 TYPE(cell_type), POINTER :: cell
928 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
929 TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
930 atomic_grid_weight_grad_t
931
932 INTEGER :: feature_begin, feature_end, feature_pos, &
933 i, iatom, j, jatom, k, local_row, &
934 natom, row
935 INTEGER, DIMENSION(2, 3) :: bo
936 LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
937 REAL(kind=dp) :: base_weight, weight_grad
938 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: weights
939 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
940 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
941 REAL(kind=dp), DIMENSION(3) :: grid_point
942 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
943
944 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
945 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
946 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
947 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
948 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
949
950 natom = SIZE(particle_set)
951 cpassert(SIZE(atom_force, 1) == 3)
952 cpassert(SIZE(atom_force, 2) == natom)
953 ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
954 dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
955 DO iatom = 1, natom
956 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
957 END DO
958
959 bo = rho_r(1)%pw_grid%bounds_local
960 local_row = 0
961 DO k = bo(1, 3), bo(2, 3)
962 DO j = bo(1, 2), bo(2, 2)
963 DO i = bo(1, 1), bo(2, 1)
964 local_row = local_row + 1
965 grid_point = native_grid_coordinate(rho_r(1)%pw_grid, [i, j, k])
966 CALL skala_gpw_smooth_partition_derivatives(grid_point, atom_coords_pbc, cell, &
967 weights, included, dweights_datom, &
968 dweights_dstrain)
969 feature_begin = features%local_feature_offsets(local_row)
970 feature_end = features%local_feature_offsets(local_row + 1) - 1
971 cpassert(feature_end - feature_begin + 1 == count(included))
972 base_weight = 0.0_dp
973 DO feature_pos = feature_begin, feature_end
974 row = features%local_feature_rows(feature_pos)
975 base_weight = base_weight + features%grid_weights(row)
976 END DO
977 feature_pos = feature_begin
978 DO iatom = 1, natom
979 IF (.NOT. included(iatom)) cycle
980 row = features%local_feature_rows(feature_pos)
981 weight_grad = grid_weight_grad(row)
982 DO jatom = 1, natom
983 atom_force(:, jatom) = atom_force(:, jatom) + &
984 weight_grad*base_weight* &
985 dweights_datom(:, jatom, iatom)
986 END DO
987 feature_pos = feature_pos + 1
988 END DO
989 cpassert(feature_pos == feature_end + 1)
990 END DO
991 END DO
992 END DO
993 cpassert(local_row == features%nflat_local)
994
995 DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
996 CALL torch_tensor_release(grid_weight_grad_t)
997 CALL torch_tensor_release(atomic_grid_weight_grad_t)
998
999 END SUBROUTINE add_smooth_partition_force
1000
1001! **************************************************************************************************
1002!> \brief Add the virial from SMOOTH native-grid atom partition weights.
1003!> \param virial_xc ...
1004!> \param features ...
1005!> \param particle_set ...
1006!> \param cell ...
1007!> \param rho_r ...
1008!> \param grid_weight_grad_t ...
1009!> \param atomic_grid_weight_grad_t ...
1010! **************************************************************************************************
1011 SUBROUTINE build_smooth_partition_virial(virial_xc, features, particle_set, cell, rho_r, &
1012 grid_weight_grad_t, atomic_grid_weight_grad_t)
1013 REAL(kind=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
1014 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1015 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
1016 TYPE(cell_type), POINTER :: cell
1017 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
1018 TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
1019 atomic_grid_weight_grad_t
1020
1021 INTEGER :: feature_begin, feature_end, feature_pos, &
1022 i, iatom, idir, j, jdir, k, local_row, &
1023 natom, row
1024 INTEGER, DIMENSION(2, 3) :: bo
1025 LOGICAL, ALLOCATABLE, DIMENSION(:) :: included
1026 REAL(kind=dp) :: base_weight, tmp, weight_grad
1027 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: weights
1028 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc
1029 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: dweights_datom, dweights_dstrain
1030 REAL(kind=dp), DIMENSION(3) :: grid_point
1031 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1032
1033 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1034 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1035 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1036 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1037 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1038
1039 natom = SIZE(particle_set)
1040 ALLOCATE (atom_coords_pbc(3, natom), included(natom), weights(natom), &
1041 dweights_datom(3, natom, natom), dweights_dstrain(3, 3, natom))
1042 DO iatom = 1, natom
1043 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
1044 END DO
1045
1046 bo = rho_r(1)%pw_grid%bounds_local
1047 local_row = 0
1048 DO k = bo(1, 3), bo(2, 3)
1049 DO j = bo(1, 2), bo(2, 2)
1050 DO i = bo(1, 1), bo(2, 1)
1051 local_row = local_row + 1
1052 grid_point = native_grid_coordinate(rho_r(1)%pw_grid, [i, j, k])
1053 CALL skala_gpw_smooth_partition_derivatives(grid_point, atom_coords_pbc, cell, &
1054 weights, included, dweights_datom, &
1055 dweights_dstrain)
1056 feature_begin = features%local_feature_offsets(local_row)
1057 feature_end = features%local_feature_offsets(local_row + 1) - 1
1058 cpassert(feature_end - feature_begin + 1 == count(included))
1059 base_weight = 0.0_dp
1060 DO feature_pos = feature_begin, feature_end
1061 row = features%local_feature_rows(feature_pos)
1062 base_weight = base_weight + features%grid_weights(row)
1063 END DO
1064 feature_pos = feature_begin
1065 DO iatom = 1, natom
1066 IF (.NOT. included(iatom)) cycle
1067 row = features%local_feature_rows(feature_pos)
1068 weight_grad = grid_weight_grad(row)
1069 DO idir = 1, 3
1070 DO jdir = 1, idir
1071 tmp = weight_grad*base_weight*dweights_dstrain(idir, jdir, iatom)
1072 virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1073 IF (idir /= jdir) virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1074 END DO
1075 END DO
1076 feature_pos = feature_pos + 1
1077 END DO
1078 cpassert(feature_pos == feature_end + 1)
1079 END DO
1080 END DO
1081 END DO
1082 cpassert(local_row == features%nflat_local)
1083
1084 DEALLOCATE (atom_coords_pbc, dweights_datom, dweights_dstrain, included, weights)
1085 CALL torch_tensor_release(grid_weight_grad_t)
1086 CALL torch_tensor_release(atomic_grid_weight_grad_t)
1087
1088 END SUBROUTINE build_smooth_partition_virial
1089
1090! **************************************************************************************************
1091!> \brief Return the Cartesian coordinate of a regular GPW grid point.
1092!> \param pw_grid ...
1093!> \param index ...
1094!> \return ...
1095! **************************************************************************************************
1096 FUNCTION native_grid_coordinate(pw_grid, index) RESULT(coord)
1097 TYPE(pw_grid_type), POINTER :: pw_grid
1098 INTEGER, DIMENSION(3), INTENT(IN) :: index
1099 REAL(kind=dp), DIMENSION(3) :: coord
1100
1101 INTEGER, DIMENSION(3) :: relative_index
1102
1103 relative_index = index - pw_grid%bounds(1, :)
1104 coord = real(relative_index(1), kind=dp)*pw_grid%dh(:, 1) + &
1105 REAL(relative_index(2), kind=dp)*pw_grid%dh(:, 2) + &
1106 REAL(relative_index(3), kind=dp)*pw_grid%dh(:, 3)
1107
1108 END FUNCTION native_grid_coordinate
1109
1110! **************************************************************************************************
1111!> \brief Evaluate a rank-local atom chunk as multiple atom-contiguous Torch subchunks.
1112!> \param features ...
1113!> \param group ...
1114!> \param max_rows ...
1115!> \param compute_grads ...
1116!> \param exc ...
1117!> \param density_grad ...
1118!> \param grad_grad ...
1119!> \param kin_grad ...
1120!> \param collapse_spin_grads ...
1121! **************************************************************************************************
1122 SUBROUTINE evaluate_atom_subchunks(features, group, max_rows, compute_grads, exc, &
1123 density_grad, grad_grad, kin_grad, collapse_spin_grads)
1124 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1125
1126 CLASS(mp_comm_type), INTENT(IN) :: group
1127 INTEGER, INTENT(IN) :: max_rows
1128 LOGICAL, INTENT(IN) :: compute_grads, collapse_spin_grads
1129 REAL(kind=dp), INTENT(OUT) :: exc
1130 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
1131 INTENT(OUT) :: density_grad, kin_grad
1132 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :), &
1133 INTENT(OUT) :: grad_grad
1134
1135 INTEGER :: base, isubchunk, local_row, nflat_local, &
1136 nroute_grad_per_point, nroute_points, &
1137 nsubchunks, phase_handle, point_pos, &
1138 subphase_handle
1139 INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
1140 route_grad_return_recv_displs, &
1141 route_grad_return_send_counts, &
1142 route_grad_return_send_displs
1143 REAL(kind=dp) :: subchunk_exc
1144 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: recv_grad_buffer, send_grad_buffer
1145 TYPE(skala_gpw_feature_type) :: subchunk
1146 TYPE(torch_tensor_type) :: subchunk_exc_tensor
1147
1148 cpassert(features%uses_atom_chunks)
1149 cpassert(max_rows > 0)
1150 nflat_local = features%nflat_local
1151 nsubchunks = skala_gpw_atom_subchunk_count(max_rows)
1152
1153 exc = 0.0_dp
1154 IF (compute_grads) THEN
1155 cpassert(features%uses_atom_chunk_routing)
1156 cpassert(sum(features%route_point_recv_counts) == features%chunk_feature_count)
1157 nroute_points = SIZE(features%route_send_local_rows)
1158 cpassert(sum(features%route_point_send_counts) == nroute_points)
1159 nroute_grad_per_point = ngrad_per_point
1160 IF (collapse_spin_grads) nroute_grad_per_point = ncollapsed_grad_per_point
1161 ALLOCATE (send_grad_buffer(max(1, nroute_grad_per_point*features%chunk_feature_count)), &
1162 recv_grad_buffer(max(1, nroute_grad_per_point*nroute_points)), &
1163 route_grad_return_send_counts(SIZE(features%route_point_recv_counts)), &
1164 route_grad_return_send_displs(SIZE(features%route_point_recv_displs)), &
1165 route_grad_return_recv_counts(SIZE(features%route_point_send_counts)), &
1166 route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
1167 route_grad_return_send_counts(:) = &
1168 nroute_grad_per_point*features%route_point_recv_counts
1169 route_grad_return_send_displs(:) = &
1170 nroute_grad_per_point*features%route_point_recv_displs
1171 route_grad_return_recv_counts(:) = &
1172 nroute_grad_per_point*features%route_point_send_counts
1173 route_grad_return_recv_displs(:) = &
1174 nroute_grad_per_point*features%route_point_send_displs
1175 END IF
1176
1177 CALL timeset("skala_gpw_atom_subchunks", phase_handle)
1178 DO isubchunk = 1, nsubchunks
1179 CALL timeset("skala_gpw_atom_subchunk_build", subphase_handle)
1180 CALL skala_gpw_feature_build_atom_subchunk(features, subchunk, isubchunk, &
1181 max_rows, compute_grads)
1182 CALL timestop(subphase_handle)
1183 CALL timeset("skala_gpw_atom_subchunk_forward", subphase_handle)
1184 CALL skala_torch_model_get_exc(cached_model, subchunk%inputs, &
1185 subchunk%grid_weights_t, subchunk_exc_tensor, &
1186 subchunk_exc)
1187 CALL timestop(subphase_handle)
1188 exc = exc + subchunk_exc
1189 IF (compute_grads) THEN
1190 CALL timeset("skala_gpw_atom_subchunk_backward", subphase_handle)
1191 CALL torch_tensor_backward_scalar(subchunk_exc_tensor)
1192 CALL timestop(subphase_handle)
1193 END IF
1194 CALL timeset("skala_gpw_atom_subchunk_release", subphase_handle)
1195 CALL torch_tensor_release(subchunk_exc_tensor)
1196 CALL skala_gpw_feature_release(subchunk)
1197 CALL timestop(subphase_handle)
1198 END DO
1199 IF (compute_grads .AND. features%chunk_feature_count > 0) THEN
1200 CALL timeset("skala_gpw_atom_subchunk_grad_pack", subphase_handle)
1201 CALL pack_atom_chunk_grads(features, send_grad_buffer, .true., collapse_spin_grads)
1202 CALL timestop(subphase_handle)
1203 END IF
1204 CALL timestop(phase_handle)
1205
1206 IF (compute_grads) THEN
1207 CALL timeset("skala_gpw_grad_route_comm", phase_handle)
1208 CALL group%alltoall(send_grad_buffer, route_grad_return_send_counts, &
1209 route_grad_return_send_displs, recv_grad_buffer, &
1210 route_grad_return_recv_counts, route_grad_return_recv_displs)
1211 CALL timestop(phase_handle)
1212
1213 CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
1214 ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
1215 kin_grad(nflat_local, 2))
1216 density_grad = 0.0_dp
1217 grad_grad = 0.0_dp
1218 kin_grad = 0.0_dp
1219 DO point_pos = 1, nroute_points
1220 local_row = features%route_send_local_rows(point_pos)
1221 cpassert(local_row >= 1 .AND. local_row <= nflat_local)
1222 base = nroute_grad_per_point*(point_pos - 1)
1223 IF (collapse_spin_grads) THEN
1224 density_grad(local_row, :) = density_grad(local_row, :) + &
1225 recv_grad_buffer(base + 1)
1226 grad_grad(local_row, 1, :) = grad_grad(local_row, 1, :) + &
1227 recv_grad_buffer(base + 2)
1228 grad_grad(local_row, 2, :) = grad_grad(local_row, 2, :) + &
1229 recv_grad_buffer(base + 3)
1230 grad_grad(local_row, 3, :) = grad_grad(local_row, 3, :) + &
1231 recv_grad_buffer(base + 4)
1232 kin_grad(local_row, :) = kin_grad(local_row, :) + recv_grad_buffer(base + 5)
1233 ELSE
1234 density_grad(local_row, :) = density_grad(local_row, :) + &
1235 recv_grad_buffer(base + 1:base + 2)
1236 grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
1237 recv_grad_buffer(base + 3)
1238 grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
1239 recv_grad_buffer(base + 4)
1240 grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
1241 recv_grad_buffer(base + 5)
1242 grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
1243 recv_grad_buffer(base + 6)
1244 grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
1245 recv_grad_buffer(base + 7)
1246 grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
1247 recv_grad_buffer(base + 8)
1248 kin_grad(local_row, :) = kin_grad(local_row, :) + &
1249 recv_grad_buffer(base + 9:base + 10)
1250 END IF
1251 END DO
1252 CALL timestop(phase_handle)
1253
1254 DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
1255 route_grad_return_recv_displs, route_grad_return_send_counts, &
1256 route_grad_return_send_displs, send_grad_buffer)
1257 END IF
1258
1259 END SUBROUTINE evaluate_atom_subchunks
1260
1261! **************************************************************************************************
1262!> \brief Select an automatic CUDA atom-subchunk row cap.
1263!> \param features ...
1264!> \param group ...
1265!> \return ...
1266! **************************************************************************************************
1267 FUNCTION auto_atom_chunk_max_rows(features, group) RESULT(max_rows)
1268 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1269
1270 CLASS(mp_comm_type), INTENT(IN) :: group
1271 INTEGER :: max_rows
1272
1273 INTEGER :: local_rows_max, target_rows
1274
1275 local_rows_max = features%chunk_feature_count
1276 CALL group%max(local_rows_max)
1277 IF (local_rows_max <= 0) THEN
1278 max_rows = 0
1279 RETURN
1280 END IF
1281
1282 IF (group%num_pe > 1) THEN
1283 target_rows = ceiling(real(local_rows_max, kind=dp)/2.0_dp)
1284 max_rows = atom_chunk_auto_row_quantum* &
1285 ((target_rows + atom_chunk_auto_row_quantum - 1)/atom_chunk_auto_row_quantum)
1286 ELSE
1287 target_rows = nint(real(local_rows_max, kind=dp)/4.0_dp)
1288 max_rows = atom_chunk_auto_row_quantum* &
1289 max(1, nint(real(target_rows, kind=dp)/ &
1290 REAL(atom_chunk_auto_row_quantum, kind=dp)))
1291 END IF
1292 max_rows = max(atom_chunk_auto_min_rows, min(atom_chunk_auto_max_rows, max_rows))
1293
1294 END FUNCTION auto_atom_chunk_max_rows
1295
1296! **************************************************************************************************
1297!> \brief Map full Torch feature gradients back to this rank's local grid order.
1298!> \param features ...
1299!> \param density_grad ...
1300!> \param grad_grad ...
1301!> \param kin_grad ...
1302! **************************************************************************************************
1303 SUBROUTINE fetch_local_feature_grads(features, density_grad, grad_grad, kin_grad)
1304 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1305 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
1306 INTENT(OUT) :: density_grad
1307 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :), &
1308 INTENT(OUT) :: grad_grad
1309 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
1310 INTENT(OUT) :: kin_grad
1311
1312 INTEGER :: feature_pos, i, j, k, local_row, row
1313 REAL(kind=dp), DIMENSION(:, :), POINTER :: density_grad_all, kin_grad_all
1314 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: grad_grad_all
1315 TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
1316
1317 NULLIFY (density_grad_all, grad_grad_all, kin_grad_all)
1318 CALL get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
1319 density_grad_all, grad_grad_all, kin_grad_all)
1320 cpassert(SIZE(density_grad_all, 1) == features%nflat)
1321 cpassert(SIZE(density_grad_all, 2) == 2)
1322 cpassert(SIZE(grad_grad_all, 1) == features%nflat)
1323 cpassert(SIZE(grad_grad_all, 2) == 3)
1324 cpassert(SIZE(grad_grad_all, 3) == 2)
1325 cpassert(SIZE(kin_grad_all, 1) == features%nflat)
1326 cpassert(SIZE(kin_grad_all, 2) == 2)
1327
1328 ALLOCATE (density_grad(features%nflat_local, 2), &
1329 grad_grad(features%nflat_local, 3, 2), &
1330 kin_grad(features%nflat_local, 2))
1331 density_grad = 0.0_dp
1332 grad_grad = 0.0_dp
1333 kin_grad = 0.0_dp
1334 local_row = 0
1335 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
1336 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
1337 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
1338 local_row = local_row + 1
1339 DO feature_pos = features%local_feature_offsets(local_row), &
1340 features%local_feature_offsets(local_row + 1) - 1
1341 row = features%local_feature_rows(feature_pos)
1342 cpassert(row >= 1 .AND. row <= features%nflat)
1343 density_grad(local_row, :) = density_grad(local_row, :) + &
1344 density_grad_all(row, :)
1345 grad_grad(local_row, :, :) = grad_grad(local_row, :, :) + &
1346 grad_grad_all(row, :, :)
1347 kin_grad(local_row, :) = kin_grad(local_row, :) + kin_grad_all(row, :)
1348 END DO
1349 END DO
1350 END DO
1351 END DO
1352 cpassert(local_row == features%nflat_local)
1353
1354 CALL torch_tensor_release(density_grad_t)
1355 CALL torch_tensor_release(grad_grad_t)
1356 CALL torch_tensor_release(kin_grad_t)
1357
1358 END SUBROUTINE fetch_local_feature_grads
1359
1360! **************************************************************************************************
1361!> \brief Pack atom-chunk Torch gradients into CP2K communication buffers.
1362!> \param features ...
1363!> \param TARGET ...
1364!> \param route_to_return_positions ...
1365!> \param collapse_spin_grads ...
1366! **************************************************************************************************
1367 SUBROUTINE pack_atom_chunk_grads(features, TARGET, route_to_return_positions, &
1368 collapse_spin_grads)
1369 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1370 REAL(kind=dp), ALLOCATABLE, DIMENSION(:), &
1371 INTENT(INOUT) :: target
1372 LOGICAL, INTENT(IN) :: route_to_return_positions
1373 LOGICAL, INTENT(IN), OPTIONAL :: collapse_spin_grads
1374
1375 INTEGER :: base, irow, ngrad_buffer_per_point, &
1376 point_pos, target_points
1377 LOGICAL :: my_collapse_spin_grads
1378 REAL(kind=dp), DIMENSION(:, :), POINTER :: chunk_density_grad, chunk_kin_grad
1379 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: chunk_grad_grad
1380 TYPE(torch_tensor_type) :: density_grad_t, grad_grad_t, kin_grad_t
1381
1382 my_collapse_spin_grads = .false.
1383 IF (PRESENT(collapse_spin_grads)) my_collapse_spin_grads = collapse_spin_grads
1384 ngrad_buffer_per_point = ngrad_per_point
1385 IF (my_collapse_spin_grads) ngrad_buffer_per_point = ncollapsed_grad_per_point
1386
1387 NULLIFY (chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
1388 CALL get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
1389 chunk_density_grad, chunk_grad_grad, chunk_kin_grad)
1390 cpassert(mod(SIZE(TARGET), ngrad_buffer_per_point) == 0)
1391 target_points = SIZE(TARGET)/ngrad_buffer_per_point
1392 cpassert(target_points >= features%chunk_feature_count)
1393 cpassert(SIZE(chunk_density_grad, 1) == features%chunk_feature_count)
1394 cpassert(SIZE(chunk_grad_grad, 1) == features%chunk_feature_count)
1395 cpassert(SIZE(chunk_grad_grad, 2) == 3)
1396 cpassert(SIZE(chunk_kin_grad, 1) == features%chunk_feature_count)
1397 IF (features%uses_collapsed_rks_dynamic) THEN
1398 cpassert(my_collapse_spin_grads)
1399 cpassert(SIZE(chunk_density_grad, 2) == 1)
1400 cpassert(SIZE(chunk_grad_grad, 3) == 1)
1401 cpassert(SIZE(chunk_kin_grad, 2) == 1)
1402 ELSE
1403 cpassert(SIZE(chunk_density_grad, 2) == 2)
1404 cpassert(SIZE(chunk_grad_grad, 3) == 2)
1405 cpassert(SIZE(chunk_kin_grad, 2) == 2)
1406 END IF
1407
1408 DO irow = 1, features%chunk_feature_count
1409 IF (route_to_return_positions) THEN
1410 point_pos = features%chunk_return_positions(irow)
1411 cpassert(point_pos >= 1 .AND. point_pos <= target_points)
1412 ELSE
1413 point_pos = irow
1414 END IF
1415 base = ngrad_buffer_per_point*(point_pos - 1)
1416 IF (my_collapse_spin_grads) THEN
1417 IF (features%uses_collapsed_rks_dynamic) THEN
1418 target(base + 1) = 0.5_dp*chunk_density_grad(irow, 1)
1419 target(base + 2) = 0.5_dp*chunk_grad_grad(irow, 1, 1)
1420 target(base + 3) = 0.5_dp*chunk_grad_grad(irow, 2, 1)
1421 target(base + 4) = 0.5_dp*chunk_grad_grad(irow, 3, 1)
1422 target(base + 5) = 0.5_dp*chunk_kin_grad(irow, 1)
1423 ELSE
1424 target(base + 1) = 0.5_dp*(chunk_density_grad(irow, 1) + &
1425 chunk_density_grad(irow, 2))
1426 target(base + 2) = 0.5_dp*(chunk_grad_grad(irow, 1, 1) + &
1427 chunk_grad_grad(irow, 1, 2))
1428 target(base + 3) = 0.5_dp*(chunk_grad_grad(irow, 2, 1) + &
1429 chunk_grad_grad(irow, 2, 2))
1430 target(base + 4) = 0.5_dp*(chunk_grad_grad(irow, 3, 1) + &
1431 chunk_grad_grad(irow, 3, 2))
1432 target(base + 5) = 0.5_dp*(chunk_kin_grad(irow, 1) + chunk_kin_grad(irow, 2))
1433 END IF
1434 ELSE
1435 target(base + 1:base + 2) = chunk_density_grad(irow, :)
1436 target(base + 3) = chunk_grad_grad(irow, 1, 1)
1437 target(base + 4) = chunk_grad_grad(irow, 2, 1)
1438 target(base + 5) = chunk_grad_grad(irow, 3, 1)
1439 target(base + 6) = chunk_grad_grad(irow, 1, 2)
1440 target(base + 7) = chunk_grad_grad(irow, 2, 2)
1441 target(base + 8) = chunk_grad_grad(irow, 3, 2)
1442 target(base + 9:base + 10) = chunk_kin_grad(irow, :)
1443 END IF
1444 END DO
1445
1446 CALL torch_tensor_release(density_grad_t)
1447 CALL torch_tensor_release(grad_grad_t)
1448 CALL torch_tensor_release(kin_grad_t)
1449
1450 END SUBROUTINE pack_atom_chunk_grads
1451
1452! **************************************************************************************************
1453!> \brief Return CPU views of autograd outputs for the SKALA dynamic feature tensors.
1454!> \param features ...
1455!> \param density_grad_t ...
1456!> \param grad_grad_t ...
1457!> \param kin_grad_t ...
1458!> \param density_grad ...
1459!> \param grad_grad ...
1460!> \param kin_grad ...
1461! **************************************************************************************************
1462 SUBROUTINE get_feature_grad_views(features, density_grad_t, grad_grad_t, kin_grad_t, &
1463 density_grad, grad_grad, kin_grad)
1464 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1465 TYPE(torch_tensor_type), INTENT(INOUT) :: density_grad_t, grad_grad_t, kin_grad_t
1466 REAL(kind=dp), DIMENSION(:, :), POINTER :: density_grad
1467 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: grad_grad
1468 REAL(kind=dp), DIMENSION(:, :), POINTER :: kin_grad
1469
1470 NULLIFY (density_grad, grad_grad, kin_grad)
1471 CALL torch_tensor_grad(features%density_t, density_grad_t)
1472 CALL torch_tensor_grad(features%grad_t, grad_grad_t)
1473 CALL torch_tensor_grad(features%kin_t, kin_grad_t)
1474 CALL torch_tensor_data_ptr(density_grad_t, density_grad)
1475 CALL torch_tensor_data_ptr(grad_grad_t, grad_grad)
1476 CALL torch_tensor_data_ptr(kin_grad_t, kin_grad)
1477
1478 END SUBROUTINE get_feature_grad_views
1479
1480! **************************************************************************************************
1481!> \brief Fetch atom-chunk gradients and route them back to their local grid owners.
1482!> \param features ...
1483!> \param group ...
1484!> \param density_grad ...
1485!> \param grad_grad ...
1486!> \param kin_grad ...
1487! **************************************************************************************************
1488 SUBROUTINE fetch_and_gather_atom_chunk_grads(features, group, density_grad, grad_grad, &
1489 kin_grad)
1490 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1491
1492 CLASS(mp_comm_type), INTENT(IN) :: group
1493 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
1494 INTENT(OUT) :: density_grad, kin_grad
1495 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :), &
1496 INTENT(OUT) :: grad_grad
1497
1498 INTEGER :: base, feature_pos, i, j, k, local_row, &
1499 nflat_local, nroute_grad_per_point, &
1500 nroute_points, phase_handle, point_pos, row
1501 INTEGER, ALLOCATABLE, DIMENSION(:) :: route_grad_return_recv_counts, &
1502 route_grad_return_recv_displs, &
1503 route_grad_return_send_counts, &
1504 route_grad_return_send_displs
1505 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: chunk_grad_buffer, global_grad_buffer, &
1506 recv_grad_buffer, send_grad_buffer
1507
1508 cpassert(features%uses_atom_chunks)
1509
1510 nflat_local = features%nflat_local
1511 IF (features%uses_atom_chunk_routing) THEN
1512 cpassert(sum(features%route_point_recv_counts) == features%chunk_feature_count)
1513 nroute_points = SIZE(features%route_send_local_rows)
1514 cpassert(sum(features%route_point_send_counts) == nroute_points)
1515
1516 nroute_grad_per_point = ngrad_per_point
1517 IF (features%uses_collapsed_rks_dynamic) THEN
1518 nroute_grad_per_point = ncollapsed_grad_per_point
1519 END IF
1520 ALLOCATE (send_grad_buffer(max(1, nroute_grad_per_point*features%chunk_feature_count)), &
1521 recv_grad_buffer(max(1, nroute_grad_per_point*nroute_points)), &
1522 route_grad_return_send_counts(SIZE(features%route_point_recv_counts)), &
1523 route_grad_return_send_displs(SIZE(features%route_point_recv_displs)), &
1524 route_grad_return_recv_counts(SIZE(features%route_point_send_counts)), &
1525 route_grad_return_recv_displs(SIZE(features%route_point_send_displs)))
1526 route_grad_return_send_counts(:) = &
1527 nroute_grad_per_point*features%route_point_recv_counts
1528 route_grad_return_send_displs(:) = &
1529 nroute_grad_per_point*features%route_point_recv_displs
1530 route_grad_return_recv_counts(:) = &
1531 nroute_grad_per_point*features%route_point_send_counts
1532 route_grad_return_recv_displs(:) = &
1533 nroute_grad_per_point*features%route_point_send_displs
1534
1535 IF (features%chunk_feature_count > 0) THEN
1536 CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
1537 CALL pack_atom_chunk_grads(features, send_grad_buffer, .true., &
1538 features%uses_collapsed_rks_dynamic)
1539 CALL timestop(phase_handle)
1540 END IF
1541
1542 CALL timeset("skala_gpw_grad_route_comm", phase_handle)
1543 CALL group%alltoall(send_grad_buffer, route_grad_return_send_counts, &
1544 route_grad_return_send_displs, recv_grad_buffer, &
1545 route_grad_return_recv_counts, route_grad_return_recv_displs)
1546 CALL timestop(phase_handle)
1547
1548 CALL timeset("skala_gpw_grad_route_scatter", phase_handle)
1549 ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
1550 kin_grad(nflat_local, 2))
1551 density_grad = 0.0_dp
1552 grad_grad = 0.0_dp
1553 kin_grad = 0.0_dp
1554 DO point_pos = 1, nroute_points
1555 local_row = features%route_send_local_rows(point_pos)
1556 cpassert(local_row >= 1 .AND. local_row <= nflat_local)
1557 base = nroute_grad_per_point*(point_pos - 1)
1558 IF (features%uses_collapsed_rks_dynamic) THEN
1559 density_grad(local_row, :) = density_grad(local_row, :) + &
1560 recv_grad_buffer(base + 1)
1561 grad_grad(local_row, 1, :) = grad_grad(local_row, 1, :) + &
1562 recv_grad_buffer(base + 2)
1563 grad_grad(local_row, 2, :) = grad_grad(local_row, 2, :) + &
1564 recv_grad_buffer(base + 3)
1565 grad_grad(local_row, 3, :) = grad_grad(local_row, 3, :) + &
1566 recv_grad_buffer(base + 4)
1567 kin_grad(local_row, :) = kin_grad(local_row, :) + recv_grad_buffer(base + 5)
1568 ELSE
1569 density_grad(local_row, :) = density_grad(local_row, :) + &
1570 recv_grad_buffer(base + 1:base + 2)
1571 grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
1572 recv_grad_buffer(base + 3)
1573 grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
1574 recv_grad_buffer(base + 4)
1575 grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
1576 recv_grad_buffer(base + 5)
1577 grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
1578 recv_grad_buffer(base + 6)
1579 grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
1580 recv_grad_buffer(base + 7)
1581 grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
1582 recv_grad_buffer(base + 8)
1583 kin_grad(local_row, :) = kin_grad(local_row, :) + &
1584 recv_grad_buffer(base + 9:base + 10)
1585 END IF
1586 END DO
1587 CALL timestop(phase_handle)
1588
1589 DEALLOCATE (recv_grad_buffer, route_grad_return_recv_counts, &
1590 route_grad_return_recv_displs, route_grad_return_send_counts, &
1591 route_grad_return_send_displs, send_grad_buffer)
1592 ELSE
1593 ALLOCATE (chunk_grad_buffer(max(1, ngrad_per_point*features%chunk_feature_count)), &
1594 global_grad_buffer(ngrad_per_point*features%nflat))
1595 IF (features%chunk_feature_count > 0) THEN
1596 CALL timeset("skala_gpw_grad_torch_pack", phase_handle)
1597 CALL pack_atom_chunk_grads(features, chunk_grad_buffer, .false.)
1598 CALL timestop(phase_handle)
1599 END IF
1600
1601 CALL timeset("skala_gpw_grad_allgatherv", phase_handle)
1602 CALL group%allgatherv(chunk_grad_buffer, global_grad_buffer, &
1603 features%chunk_grad_counts, features%chunk_grad_displs)
1604 CALL timestop(phase_handle)
1605
1606 CALL timeset("skala_gpw_grad_scatter", phase_handle)
1607 ALLOCATE (density_grad(nflat_local, 2), grad_grad(nflat_local, 3, 2), &
1608 kin_grad(nflat_local, 2))
1609 density_grad = 0.0_dp
1610 grad_grad = 0.0_dp
1611 kin_grad = 0.0_dp
1612 local_row = 0
1613 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
1614 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
1615 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
1616 local_row = local_row + 1
1617 DO feature_pos = features%local_feature_offsets(local_row), &
1618 features%local_feature_offsets(local_row + 1) - 1
1619 row = features%local_feature_rows(feature_pos)
1620 cpassert(row >= 1 .AND. row <= features%nflat)
1621 base = ngrad_per_point*(row - 1)
1622 density_grad(local_row, :) = density_grad(local_row, :) + &
1623 global_grad_buffer(base + 1:base + 2)
1624 grad_grad(local_row, 1, 1) = grad_grad(local_row, 1, 1) + &
1625 global_grad_buffer(base + 3)
1626 grad_grad(local_row, 2, 1) = grad_grad(local_row, 2, 1) + &
1627 global_grad_buffer(base + 4)
1628 grad_grad(local_row, 3, 1) = grad_grad(local_row, 3, 1) + &
1629 global_grad_buffer(base + 5)
1630 grad_grad(local_row, 1, 2) = grad_grad(local_row, 1, 2) + &
1631 global_grad_buffer(base + 6)
1632 grad_grad(local_row, 2, 2) = grad_grad(local_row, 2, 2) + &
1633 global_grad_buffer(base + 7)
1634 grad_grad(local_row, 3, 2) = grad_grad(local_row, 3, 2) + &
1635 global_grad_buffer(base + 8)
1636 kin_grad(local_row, :) = kin_grad(local_row, :) + &
1637 global_grad_buffer(base + 9:base + 10)
1638 END DO
1639 END DO
1640 END DO
1641 END DO
1642 CALL timestop(phase_handle)
1643 DEALLOCATE (chunk_grad_buffer, global_grad_buffer)
1644
1645 END IF
1646
1647 END SUBROUTINE fetch_and_gather_atom_chunk_grads
1648
1649! **************************************************************************************************
1650!> \brief Build the native SKALA XC virial from feature gradients.
1651!> \param virial_xc ...
1652!> \param rho_set ...
1653!> \param rho_r ...
1654!> \param grad_grad ...
1655! **************************************************************************************************
1656 SUBROUTINE build_virial_from_feature_grads(virial_xc, rho_set, rho_r, grad_grad)
1657 REAL(kind=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
1658 TYPE(xc_rho_set_type), INTENT(IN) :: rho_set
1659 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: rho_r
1660 REAL(kind=dp), DIMENSION(:, :, :), INTENT(IN) :: grad_grad
1661
1662 INTEGER :: i, idir, ipt, ispin, j, jdir, k, nspins
1663 INTEGER, DIMENSION(2, 3) :: bo
1664 REAL(kind=dp) :: grad_i, tmp
1665 TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
1666
1667 nspins = SIZE(rho_r)
1668 bo = rho_r(1)%pw_grid%bounds_local
1669 ipt = 0
1670
1671 IF (nspins == 1) THEN
1672 CALL xc_rho_set_get(rho_set, drho=drho)
1673 DO k = bo(1, 3), bo(2, 3)
1674 DO j = bo(1, 2), bo(2, 2)
1675 DO i = bo(1, 1), bo(2, 1)
1676 ipt = ipt + 1
1677 DO idir = 1, 3
1678 grad_i = 0.5_dp*(grad_grad(ipt, idir, 1) + grad_grad(ipt, idir, 2))
1679 DO jdir = 1, idir
1680 tmp = -grad_i*drho(jdir)%array(i, j, k)
1681 virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1682 virial_xc(idir, jdir) = virial_xc(jdir, idir)
1683 END DO
1684 END DO
1685 END DO
1686 END DO
1687 END DO
1688 ELSE
1689 CALL xc_rho_set_get(rho_set, drhoa=drhoa, drhob=drhob)
1690 DO k = bo(1, 3), bo(2, 3)
1691 DO j = bo(1, 2), bo(2, 2)
1692 DO i = bo(1, 1), bo(2, 1)
1693 ipt = ipt + 1
1694 DO idir = 1, 3
1695 DO jdir = 1, idir
1696 tmp = 0.0_dp
1697 DO ispin = 1, 2
1698 IF (ispin == 1) THEN
1699 tmp = tmp - grad_grad(ipt, idir, ispin)*drhoa(jdir)%array(i, j, k)
1700 ELSE
1701 tmp = tmp - grad_grad(ipt, idir, ispin)*drhob(jdir)%array(i, j, k)
1702 END IF
1703 END DO
1704 virial_xc(jdir, idir) = virial_xc(jdir, idir) + tmp
1705 virial_xc(idir, jdir) = virial_xc(jdir, idir)
1706 END DO
1707 END DO
1708 END DO
1709 END DO
1710 END DO
1711 END IF
1712
1713 END SUBROUTINE build_virial_from_feature_grads
1714
1715! **************************************************************************************************
1716!> \brief Print a native SKALA XC virial contribution for diagnostics.
1717!> \param label ...
1718!> \param delta ...
1719!> \param root_rank ...
1720! **************************************************************************************************
1721 SUBROUTINE print_virial_delta(label, delta, root_rank)
1722 CHARACTER(LEN=*), INTENT(IN) :: label
1723 REAL(kind=dp), DIMENSION(3, 3), INTENT(IN) :: delta
1724 LOGICAL, INTENT(IN) :: root_rank
1725
1726 INTEGER :: i, iw
1727
1728 IF (.NOT. root_rank) RETURN
1729 iw = cp_logger_get_default_io_unit()
1730 IF (iw <= 0) RETURN
1731 WRITE (iw, "(T2,A,1X,A)") "SKALA_GPW| XC virial contribution", trim(label)
1732 DO i = 1, 3
1733 WRITE (iw, "(T2,A,1X,3ES20.10)") "SKALA_GPW|", delta(i, 1:3)
1734 END DO
1735
1736 END SUBROUTINE print_virial_delta
1737
1738! **************************************************************************************************
1739!> \brief Add explicit SKALA coordinate-feature contributions to the XC virial.
1740!> \param virial_xc ...
1741!> \param features ...
1742!> \param atom_coord_grad_t ...
1743!> \param grid_coord_grad_t ...
1744!> \param root_rank ...
1745!> \param print_components ...
1746! **************************************************************************************************
1747 SUBROUTINE build_static_coordinate_virial(virial_xc, features, atom_coord_grad_t, &
1748 grid_coord_grad_t, root_rank, print_components)
1749 REAL(kind=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
1750 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1751 TYPE(torch_tensor_type), INTENT(INOUT) :: atom_coord_grad_t, grid_coord_grad_t
1752 LOGICAL, INTENT(IN) :: root_rank
1753 LOGICAL, INTENT(IN), OPTIONAL :: print_components
1754
1755 INTEGER :: feature_pos, i, iatom, idir, iw, j, &
1756 jdir, k, local_row, row
1757 LOGICAL :: my_print_components
1758 REAL(kind=dp) :: tmp
1759 REAL(kind=dp), DIMENSION(3, 3) :: atom_virial, grid_virial
1760 REAL(kind=dp), DIMENSION(:, :), POINTER :: atom_coord_grad, grid_coord_grad
1761
1762 my_print_components = .false.
1763 IF (PRESENT(print_components)) my_print_components = print_components
1764
1765 NULLIFY (atom_coord_grad, grid_coord_grad)
1766 CALL torch_tensor_grad(features%grid_coords_t, grid_coord_grad_t)
1767 CALL torch_tensor_data_ptr(grid_coord_grad_t, grid_coord_grad)
1768 CALL torch_tensor_data_ptr(atom_coord_grad_t, atom_coord_grad)
1769
1770 grid_virial = 0.0_dp
1771 atom_virial = 0.0_dp
1772 local_row = 0
1773 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
1774 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
1775 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
1776 local_row = local_row + 1
1777 DO feature_pos = features%local_feature_offsets(local_row), &
1778 features%local_feature_offsets(local_row + 1) - 1
1779 row = features%local_feature_rows(feature_pos)
1780 DO idir = 1, 3
1781 DO jdir = 1, 3
1782 tmp = grid_coord_grad(idir, row)*features%grid_coords(jdir, row)
1783 grid_virial(idir, jdir) = grid_virial(idir, jdir) + tmp
1784 virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1785 END DO
1786 END DO
1787 END DO
1788 END DO
1789 END DO
1790 END DO
1791 cpassert(local_row == features%nflat_local)
1792
1793 IF (root_rank) THEN
1794 DO iatom = 1, SIZE(features%coarse_0_atomic_coords, 2)
1795 DO idir = 1, 3
1796 DO jdir = 1, 3
1797 tmp = atom_coord_grad(idir, iatom)*features%coarse_0_atomic_coords(jdir, iatom)
1798 atom_virial(idir, jdir) = atom_virial(idir, jdir) + tmp
1799 virial_xc(idir, jdir) = virial_xc(idir, jdir) + tmp
1800 END DO
1801 END DO
1802 END DO
1803 END IF
1804
1805 IF (my_print_components .AND. root_rank) THEN
1806 iw = cp_logger_get_default_io_unit()
1807 IF (iw > 0) THEN
1808 CALL print_virial_delta("static-grid", grid_virial, .true.)
1809 CALL print_virial_delta("static-atom", atom_virial, .true.)
1810 END IF
1811 END IF
1812
1813 CALL torch_tensor_release(grid_coord_grad_t)
1814
1815 END SUBROUTINE build_static_coordinate_virial
1816
1817! **************************************************************************************************
1818!> \brief Add residual SKALA weight-feature contributions to the XC virial.
1819!> \param virial_xc ...
1820!> \param features ...
1821!> \param exc ...
1822!> \param grid_weight_grad_t ...
1823!> \param atomic_grid_weight_grad_t ...
1824!> \param root_rank ...
1825!> \param print_components ...
1826! **************************************************************************************************
1827 SUBROUTINE build_weight_virial(virial_xc, features, exc, grid_weight_grad_t, &
1828 atomic_grid_weight_grad_t, root_rank, print_components)
1829 REAL(kind=dp), DIMENSION(3, 3), INTENT(INOUT) :: virial_xc
1830 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1831 REAL(kind=dp), INTENT(IN) :: exc
1832 TYPE(torch_tensor_type), INTENT(INOUT) :: grid_weight_grad_t, &
1833 atomic_grid_weight_grad_t
1834 LOGICAL, INTENT(IN) :: root_rank
1835 LOGICAL, INTENT(IN), OPTIONAL :: print_components
1836
1837 INTEGER :: feature_pos, i, idir, iw, j, k, &
1838 local_row, row
1839 LOGICAL :: my_print_components
1840 REAL(kind=dp) :: atomic_tmp, exc_tmp, grid_tmp, tmp
1841 REAL(kind=dp), DIMENSION(:), POINTER :: atomic_grid_weight_grad, grid_weight_grad
1842
1843 my_print_components = .false.
1844 IF (PRESENT(print_components)) my_print_components = print_components
1845
1846 NULLIFY (atomic_grid_weight_grad, grid_weight_grad)
1847 CALL torch_tensor_grad(features%grid_weights_t, grid_weight_grad_t)
1848 CALL torch_tensor_grad(features%atomic_grid_weights_t, atomic_grid_weight_grad_t)
1849 CALL torch_tensor_data_ptr(grid_weight_grad_t, grid_weight_grad)
1850 CALL torch_tensor_data_ptr(atomic_grid_weight_grad_t, atomic_grid_weight_grad)
1851
1852 grid_tmp = 0.0_dp
1853 atomic_tmp = 0.0_dp
1854 local_row = 0
1855 DO k = lbound(features%feature_index, 3), ubound(features%feature_index, 3)
1856 DO j = lbound(features%feature_index, 2), ubound(features%feature_index, 2)
1857 DO i = lbound(features%feature_index, 1), ubound(features%feature_index, 1)
1858 local_row = local_row + 1
1859 DO feature_pos = features%local_feature_offsets(local_row), &
1860 features%local_feature_offsets(local_row + 1) - 1
1861 row = features%local_feature_rows(feature_pos)
1862 grid_tmp = grid_tmp + grid_weight_grad(row)*features%grid_weights(row)
1863 atomic_tmp = atomic_tmp + &
1864 atomic_grid_weight_grad(row)*features%atomic_grid_weights(row)
1865 END DO
1866 END DO
1867 END DO
1868 END DO
1869 cpassert(local_row == features%nflat_local)
1870 exc_tmp = 0.0_dp
1871 IF (root_rank) exc_tmp = -exc
1872 tmp = grid_tmp + atomic_tmp + exc_tmp
1873
1874 IF (my_print_components .AND. root_rank) THEN
1875 iw = cp_logger_get_default_io_unit()
1876 IF (iw > 0) THEN
1877 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight grid", grid_tmp
1878 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight atomic", atomic_tmp
1879 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight final", exc_tmp
1880 WRITE (iw, "(T2,A,1X,ES20.10)") "SKALA_GPW| XC virial weight residual", tmp
1881 END IF
1882 END IF
1883
1884 DO idir = 1, 3
1885 virial_xc(idir, idir) = virial_xc(idir, idir) + tmp
1886 END DO
1887
1888 CALL torch_tensor_release(grid_weight_grad_t)
1889 CALL torch_tensor_release(atomic_grid_weight_grad_t)
1890
1891 END SUBROUTINE build_weight_virial
1892
1893! **************************************************************************************************
1894!> \brief Fill CP2K VXC real-space arrays from Torch feature gradients.
1895!> \param vxc_rho ...
1896!> \param vxc_tau ...
1897!> \param rho_r ...
1898!> \param pw_pool ...
1899!> \param density_grad ...
1900!> \param grad_grad ...
1901!> \param kin_grad ...
1902!> \param xc_deriv_method_id ...
1903! **************************************************************************************************
1904 SUBROUTINE build_vxc_from_feature_grads(vxc_rho, vxc_tau, rho_r, pw_pool, &
1905 density_grad, grad_grad, kin_grad, &
1906 xc_deriv_method_id)
1907 TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER :: vxc_rho, vxc_tau, rho_r
1908 TYPE(pw_pool_type), POINTER :: pw_pool
1909 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: density_grad
1910 REAL(kind=dp), DIMENSION(:, :, :), INTENT(IN) :: grad_grad
1911 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: kin_grad
1912 INTEGER, INTENT(IN) :: xc_deriv_method_id
1913
1914 INTEGER :: i, ipt, ispin, j, k, nspins
1915 INTEGER, DIMENSION(2, 3) :: bo
1916 REAL(kind=dp) :: dvol_inv
1917 TYPE(pw_c1d_gs_type) :: tmp_g, vxc_g
1918 TYPE(pw_r3d_rs_type), DIMENSION(3) :: grad_pw
1919
1920 nspins = SIZE(rho_r)
1921 bo = rho_r(1)%pw_grid%bounds_local
1922 dvol_inv = 1.0_dp/rho_r(1)%pw_grid%dvol
1923
1924 ALLOCATE (vxc_rho(nspins), vxc_tau(nspins))
1925 DO ispin = 1, nspins
1926 CALL pw_pool%create_pw(vxc_rho(ispin))
1927 CALL pw_pool%create_pw(vxc_tau(ispin))
1928 CALL pw_zero(vxc_rho(ispin))
1929 CALL pw_zero(vxc_tau(ispin))
1930 END DO
1931
1932 IF (xc_requires_tmp_g(xc_deriv_method_id) .OR. rho_r(1)%pw_grid%spherical) THEN
1933 CALL pw_pool%create_pw(vxc_g)
1934 IF (.NOT. rho_r(1)%pw_grid%spherical) CALL pw_pool%create_pw(tmp_g)
1935 END IF
1936
1937 DO ispin = 1, nspins
1938 DO i = 1, 3
1939 CALL pw_pool%create_pw(grad_pw(i))
1940 CALL pw_zero(grad_pw(i))
1941 END DO
1942
1943 ipt = 0
1944 DO k = bo(1, 3), bo(2, 3)
1945 DO j = bo(1, 2), bo(2, 2)
1946 DO i = bo(1, 1), bo(2, 1)
1947 ipt = ipt + 1
1948 IF (nspins == 1) THEN
1949 vxc_rho(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
1950 (density_grad(ipt, 1) + density_grad(ipt, 2))
1951 vxc_tau(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
1952 (kin_grad(ipt, 1) + kin_grad(ipt, 2))
1953 grad_pw(1)%array(i, j, k) = 0.5_dp*dvol_inv* &
1954 (grad_grad(ipt, 1, 1) + grad_grad(ipt, 1, 2))
1955 grad_pw(2)%array(i, j, k) = 0.5_dp*dvol_inv* &
1956 (grad_grad(ipt, 2, 1) + grad_grad(ipt, 2, 2))
1957 grad_pw(3)%array(i, j, k) = 0.5_dp*dvol_inv* &
1958 (grad_grad(ipt, 3, 1) + grad_grad(ipt, 3, 2))
1959 ELSE
1960 vxc_rho(ispin)%array(i, j, k) = dvol_inv*density_grad(ipt, ispin)
1961 vxc_tau(ispin)%array(i, j, k) = dvol_inv*kin_grad(ipt, ispin)
1962 grad_pw(1)%array(i, j, k) = dvol_inv*grad_grad(ipt, 1, ispin)
1963 grad_pw(2)%array(i, j, k) = dvol_inv*grad_grad(ipt, 2, ispin)
1964 grad_pw(3)%array(i, j, k) = dvol_inv*grad_grad(ipt, 3, ispin)
1965 END IF
1966 END DO
1967 END DO
1968 END DO
1969
1970 DO i = 1, 3
1971 CALL pw_scale(grad_pw(i), -1.0_dp)
1972 END DO
1973 CALL xc_pw_divergence(xc_deriv_method_id, grad_pw, tmp_g, vxc_g, vxc_rho(ispin))
1974
1975 DO i = 1, 3
1976 CALL pw_pool%give_back_pw(grad_pw(i))
1977 END DO
1978 END DO
1979
1980 IF (ASSOCIATED(vxc_g%pw_grid)) CALL pw_pool%give_back_pw(vxc_g)
1981 IF (ASSOCIATED(tmp_g%pw_grid)) CALL pw_pool%give_back_pw(tmp_g)
1982
1983 END SUBROUTINE build_vxc_from_feature_grads
1984
1985! **************************************************************************************************
1986!> \brief Print optional diagnostics for the CP2K-native SKALA GPW feature block.
1987!> \param features ...
1988!> \param print_active ...
1989! **************************************************************************************************
1990 SUBROUTINE print_native_grid_diagnostics(features, print_active)
1991 TYPE(skala_gpw_feature_type), INTENT(IN) :: features
1992 LOGICAL, INTENT(IN) :: print_active
1993
1994 INTEGER :: atom_rows_max, atom_rows_min, &
1995 chunk_rows_max, chunk_rows_min, iw
1996 REAL(kind=dp) :: chunk_imbalance
1997
1998 IF (.NOT. print_active) RETURN
1999
2000 iw = cp_logger_get_default_io_unit()
2001 IF (iw <= 0) RETURN
2002 WRITE (unit=iw, fmt="(/,T2,A,1X,ES19.11)") &
2003 "SKALA_GPW| Native grid feature electrons", features%electron_count
2004 WRITE (unit=iw, fmt="(T2,A,1X,ES19.11)") &
2005 "SKALA_GPW| Native grid feature spin moment", features%spin_moment
2006 WRITE (unit=iw, fmt="(T2,A,1X,ES19.11)") &
2007 "SKALA_GPW| Native grid feature weight sum", features%grid_weight_sum
2008 IF (ALLOCATED(features%atomic_grid_sizes)) THEN
2009 atom_rows_min = int(minval(features%atomic_grid_sizes))
2010 atom_rows_max = int(maxval(features%atomic_grid_sizes))
2011 WRITE (unit=iw, fmt="(T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2012 "SKALA_GPW| Native grid atom row range", atom_rows_min, "to", &
2013 atom_rows_max, "sum", int(sum(features%atomic_grid_sizes))
2014 END IF
2015 IF (features%uses_atom_chunks) THEN
2016 WRITE (unit=iw, fmt="(T2,A,1X,I0,1X,A,1X,I0)") &
2017 "SKALA_GPW| Native grid atom chunk rows", features%chunk_feature_count, &
2018 "of", features%nflat
2019 IF (ALLOCATED(features%chunk_grad_counts)) THEN
2020 chunk_rows_min = minval(features%chunk_grad_counts)/ngrad_per_point
2021 chunk_rows_max = maxval(features%chunk_grad_counts)/ngrad_per_point
2022 chunk_imbalance = real(chunk_rows_max, kind=dp)/real(max(1, chunk_rows_min), kind=dp)
2023 WRITE (unit=iw, fmt="(T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,ES12.5)") &
2024 "SKALA_GPW| Native grid atom chunk row range", chunk_rows_min, &
2025 "to", chunk_rows_max, "imbalance", chunk_imbalance
2026 END IF
2027 END IF
2028
2029 END SUBROUTINE print_native_grid_diagnostics
2030
2031! **************************************************************************************************
2032!> \brief Configure CUDA device selection for the native SKALA GPW Torch path.
2033!> \param use_cuda ...
2034!> \param requested_device ...
2035!> \param group ...
2036!> \return selected CUDA device, or -1 for CPU fallback/no visible CUDA device
2037! **************************************************************************************************
2038 FUNCTION configure_native_grid_cuda(use_cuda, requested_device, group) RESULT(selected_device)
2039 LOGICAL, INTENT(IN) :: use_cuda
2040 INTEGER, INTENT(IN) :: requested_device
2041
2042 CLASS(mp_comm_type), INTENT(IN) :: group
2043
2044 INTEGER :: cuda_device_count, iw, pe, selected_device
2045 INTEGER, ALLOCATABLE, DIMENSION(:) :: selected_devices
2046
2047 selected_device = -1
2048
2049 IF (.NOT. use_cuda) RETURN
2050
2051 IF (.NOT. torch_cuda_is_available()) THEN
2052 cuda_device_count = 0
2053 ELSE
2054 cuda_device_count = torch_cuda_device_count()
2055 END IF
2056 IF (cuda_device_count > 0) THEN
2057 IF (requested_device < 0) THEN
2058 selected_device = mod(group%mepos, cuda_device_count)
2059 ELSE
2060 selected_device = requested_device
2061 END IF
2062 END IF
2063 IF (selected_device >= cuda_device_count) THEN
2064 CALL cp_abort(__location__, &
2065 "GAUXC%NATIVE_GRID_CUDA_DEVICE selects a CUDA device outside the visible "// &
2066 "Torch CUDA device range.")
2067 END IF
2068 IF (selected_device >= 0) CALL offload_set_chosen_device(selected_device)
2069
2070 ALLOCATE (selected_devices(group%num_pe))
2071 CALL group%allgather(selected_device, selected_devices)
2072
2073 IF (group%mepos /= 0) RETURN
2074 IF (selected_device == logged_cuda_device .AND. &
2075 cuda_device_count == logged_cuda_device_count .AND. &
2076 group%num_pe == logged_cuda_nproc .AND. &
2077 requested_device == logged_cuda_request) RETURN
2078
2079 iw = cp_logger_get_default_io_unit()
2080 IF (iw <= 0) RETURN
2081 IF (selected_device >= 0) THEN
2082 WRITE (unit=iw, fmt="(/,T2,A,1X,I0,1X,A,1X,I0,1X,A,1X,I0)") &
2083 "SKALA_GPW| Native grid Torch CUDA device", selected_device, &
2084 "of", cuda_device_count, "requested", requested_device
2085 ELSE
2086 WRITE (unit=iw, fmt="(/,T2,A)") &
2087 "SKALA_GPW| Native grid Torch CUDA requested, but no Torch CUDA device is visible"
2088 END IF
2089 WRITE (unit=iw, fmt="(T2,A)", advance="NO") &
2090 "SKALA_GPW| Native grid Torch CUDA rank devices"
2091 DO pe = 1, group%num_pe
2092 WRITE (unit=iw, fmt="(1X,I0,A,I0)", advance="NO") pe - 1, ":", selected_devices(pe)
2093 END DO
2094 WRITE (unit=iw, fmt=*)
2095
2096 logged_cuda_device = selected_device
2097 logged_cuda_device_count = cuda_device_count
2098 logged_cuda_nproc = group%num_pe
2099 logged_cuda_request = requested_device
2100
2101 END FUNCTION configure_native_grid_cuda
2102
2103! **************************************************************************************************
2104!> \brief Load and cache the TorchScript SKALA model.
2105!> \param model_path ...
2106!> \param cuda_device ...
2107! **************************************************************************************************
2108 SUBROUTINE ensure_model_loaded(model_path, cuda_device)
2109 CHARACTER(len=*), INTENT(IN) :: model_path
2110 INTEGER, INTENT(IN) :: cuda_device
2111
2112 IF (cached_model_loaded) THEN
2113 IF (trim(cached_model_path) == trim(model_path) .AND. &
2114 cached_model_cuda_device == cuda_device) RETURN
2115 CALL skala_torch_model_release(cached_model)
2116 cached_model_loaded = .false.
2117 END IF
2118
2119 CALL skala_torch_model_load(cached_model, trim(model_path))
2120 cached_model_path = model_path
2121 cached_model_cuda_device = cuda_device
2122 cached_model_loaded = .true.
2123
2124 END SUBROUTINE ensure_model_loaded
2125
2126! **************************************************************************************************
2127!> \brief Resolve the SKALA TorchScript model path from the GAUXC subsection.
2128!> \param xc_section ...
2129!> \param model_path ...
2130! **************************************************************************************************
2131 SUBROUTINE get_skala_model_path(xc_section, model_path)
2132 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
2133 CHARACTER(len=default_path_length), INTENT(OUT) :: model_path
2134
2135 CHARACTER(len=default_path_length) :: model_key
2136 INTEGER :: env_status
2137 LOGICAL :: native_grid_use_cuda
2138 TYPE(section_vals_type), POINTER :: gauxc_section
2139
2140 gauxc_section => get_gauxc_section(xc_section)
2141 IF (.NOT. ASSOCIATED(gauxc_section)) THEN
2142 cpabort("Native SKALA GPW requires an XC_FUNCTIONAL%GAUXC section")
2143 END IF
2144
2145 CALL section_vals_val_get(gauxc_section, "MODEL", c_val=model_path)
2146 model_key = adjustl(model_path)
2147 CALL uppercase(model_key)
2148 IF (trim(model_key) == "NONE" .OR. trim(model_key) == "") THEN
2149 cpabort("Native SKALA GPW requires GAUXC%MODEL SKALA or a TorchScript model path")
2150 ELSE IF (trim(model_key) == "SKALA") THEN
2151 CALL section_vals_val_get(gauxc_section, "NATIVE_GRID_USE_CUDA", l_val=native_grid_use_cuda)
2152 IF (native_grid_use_cuda) THEN
2153 CALL get_environment_variable("GAUXC_SKALA_CUDA_MODEL", model_path, status=env_status)
2154 IF (env_status == 0 .AND. len_trim(model_path) > 0) RETURN
2155 END IF
2156 CALL get_environment_variable("GAUXC_SKALA_MODEL", model_path, status=env_status)
2157 IF (env_status /= 0 .OR. len_trim(model_path) == 0) THEN
2158 IF (native_grid_use_cuda) THEN
2159 CALL cp_abort(__location__, &
2160 "MODEL SKALA CUDA path requires GAUXC_SKALA_CUDA_MODEL or GAUXC_SKALA_MODEL")
2161 ELSE
2162 CALL cp_abort(__location__, &
2163 "MODEL SKALA requires the GAUXC_SKALA_MODEL environment variable")
2164 END IF
2165 END IF
2166 END IF
2167
2168 END SUBROUTINE get_skala_model_path
2169
2170! **************************************************************************************************
2171!> \brief Return the first GAUXC functional subsection, if present.
2172!> \param xc_section ...
2173!> \return ...
2174! **************************************************************************************************
2175 FUNCTION get_gauxc_section(xc_section) RESULT(gauxc_section)
2176 TYPE(section_vals_type), INTENT(IN), POINTER :: xc_section
2177 TYPE(section_vals_type), POINTER :: gauxc_section
2178
2179 INTEGER :: ifun
2180 TYPE(section_vals_type), POINTER :: functionals, xc_fun
2181
2182 NULLIFY (gauxc_section)
2183 IF (.NOT. ASSOCIATED(xc_section)) RETURN
2184
2185 functionals => section_vals_get_subs_vals(xc_section, "XC_FUNCTIONAL")
2186 IF (.NOT. ASSOCIATED(functionals)) RETURN
2187
2188 ifun = 0
2189 DO
2190 ifun = ifun + 1
2191 xc_fun => section_vals_get_subs_vals2(functionals, i_section=ifun)
2192 IF (.NOT. ASSOCIATED(xc_fun)) EXIT
2193 IF (xc_fun%section%name == "GAUXC") THEN
2194 gauxc_section => xc_fun
2195 EXIT
2196 END IF
2197 END DO
2198
2199 END FUNCTION get_gauxc_section
2200
2201END MODULE skala_gpw_functional
cell_types::pbc
Definition cell_types.F:103
pw_methods::pw_scale
Definition pw_methods.F:94
pw_methods::pw_zero
Definition pw_methods.F:83
torch_api::torch_tensor_data_ptr
Definition torch_api.F:65
torch_api::torch_tensor_from_array
Definition torch_api.F:44
cell_types
Handles all functions related to the CELL.
Definition cell_types.F:15
cp_array_utils
various utilities that regard array of different kinds: output, allocation,... maybe it is not a good...
Definition cp_array_utils.F:18
cp_log_handling
various routines to log and control the output. The idea is that decisions about where to log should ...
Definition cp_log_handling.F:41
cp_log_handling::cp_logger_get_default_io_unit
integer function, public cp_logger_get_default_io_unit(logger)
returns the unit nr for the ionode (-1 on all other processors) skips as well checks if the procs cal...
Definition cp_log_handling.F:523
input_section_types
objects that represent the structure of input sections and the data contained in an input section
Definition input_section_types.F:15
input_section_types::section_get_rval
real(kind=dp) function, public section_get_rval(section_vals, keyword_name)
...
Definition input_section_types.F:953
input_section_types::section_vals_get_subs_vals2
type(section_vals_type) function, pointer, public section_vals_get_subs_vals2(section_vals, i_section, i_rep_section)
returns the values of the n-th non default subsection (null if no such section exists (not so many no...
Definition input_section_types.F:791
input_section_types::section_vals_get_subs_vals
recursive type(section_vals_type) function, pointer, public section_vals_get_subs_vals(section_vals, subsection_name, i_rep_section, can_return_null)
returns the values of the requested subsection
Definition input_section_types.F:735
input_section_types::section_vals_val_get
subroutine, public section_vals_val_get(section_vals, keyword_name, i_rep_section, i_rep_val, n_rep_val, val, l_val, i_val, r_val, c_val, l_vals, i_vals, r_vals, c_vals, explicit)
returns the requested value
Definition input_section_types.F:1052
kinds
Defines the basic variable types.
Definition kinds.F:23
kinds::int_8
integer, parameter, public int_8
Definition kinds.F:54
kinds::dp
integer, parameter, public dp
Definition kinds.F:34
kinds::default_path_length
integer, parameter, public default_path_length
Definition kinds.F:58
message_passing
Interface to the message passing library MPI.
Definition message_passing.F:23
offload_api
Fortran API for the offload package, which is written in C.
Definition offload_api.F:12
offload_api::offload_set_chosen_device
subroutine, public offload_set_chosen_device(device_id)
Selects the chosen device to be used.
Definition offload_api.F:129
particle_types
Define the data structure for the particle information.
Definition particle_types.F:19
pw_grid_types
Definition pw_grid_types.F:13
pw_methods
Definition pw_methods.F:25
pw_pool_types
Manages a pool of grids (to be used for example as tmp objects), but can also be used to instantiate ...
Definition pw_pool_types.F:24
pw_types
Definition pw_types.F:24
qs_grid_atom
Definition qs_grid_atom.F:8
skala_gpw_features
Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
Definition skala_gpw_features.F:11
skala_gpw_features::skala_gpw_feature_build_atom_subchunk
subroutine, public skala_gpw_feature_build_atom_subchunk(parent, features, subchunk_index, max_rows, requires_grad)
Build an atom-contiguous subchunk feature bundle from a rank-local atom chunk.
Definition skala_gpw_features.F:1694
skala_gpw_features::skala_gpw_atom_partition_hard
integer, parameter, public skala_gpw_atom_partition_hard
Definition skala_gpw_features.F:38
skala_gpw_features::skala_gpw_feature_release
subroutine, public skala_gpw_feature_release(features)
Release Torch objects and backing arrays owned by a feature bundle.
Definition skala_gpw_features.F:1553
skala_gpw_features::skala_gpw_atom_subchunk_count
integer function, public skala_gpw_atom_subchunk_count(max_rows)
Return how many atom-contiguous subchunks the cached rank chunk needs.
Definition skala_gpw_features.F:1657
skala_gpw_features::skala_gpw_feature_build
subroutine, public skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, requires_grad, weights, requires_coordinate_grad, requires_stress_grad, use_atom_chunks, route_atom_chunks, atom_partition)
Build a flat SKALA molecular feature dictionary from a local GPW grid.
Definition skala_gpw_features.F:200
skala_gpw_features::skala_gpw_smooth_partition_derivatives
subroutine, public skala_gpw_smooth_partition_derivatives(grid_point, atom_coords, cell, weights, included, dweights_datom, dweights_dstrain)
Build smooth atom weights and their atom/cell deformation derivatives.
Definition skala_gpw_features.F:2208
skala_gpw_features::skala_gpw_atom_partition_smooth
integer, parameter, public skala_gpw_atom_partition_smooth
Definition skala_gpw_features.F:38
skala_gpw_functional
Experimental CP2K-native GPW real-space-grid path for SKALA TorchScript models.
Definition skala_gpw_functional.F:11
skala_gpw_functional::skala_gapw_density_partition_soft_only
integer, parameter, public skala_gapw_density_partition_soft_only
Definition skala_gpw_functional.F:72
skala_gpw_functional::ensure_native_skala_grid_scope
subroutine, public ensure_native_skala_grid_scope(xc_section)
Enforce the currently implemented native SKALA GPW input scope.
Definition skala_gpw_functional.F:174
skala_gpw_functional::skala_gpw_eval
subroutine, public skala_gpw_eval(vxc_rho, vxc_tau, exc, rho_r, rho_g, tau, xc_section, weights, pw_pool, particle_set, cell, compute_virial, virial_xc, just_energy, atom_force)
Evaluate SKALA energy and first derivatives on a CP2K GPW grid.
Definition skala_gpw_functional.F:241
skala_gpw_functional::skala_gapw_density_partition_none
integer, parameter, public skala_gapw_density_partition_none
Definition skala_gpw_functional.F:72
skala_gpw_functional::xc_section_uses_gauxc_model
logical function, public xc_section_uses_gauxc_model(xc_section)
Return true if the GAUXC subsection requests a model evaluation.
Definition skala_gpw_functional.F:117
skala_gpw_functional::skala_gapw_density_partition_hard_minus_soft
integer, parameter, public skala_gapw_density_partition_hard_minus_soft
Definition skala_gpw_functional.F:72
skala_gpw_functional::native_skala_gapw_density_partition
integer function, public native_skala_gapw_density_partition(xc_section)
Return the hard/soft GAPW one-center density partition for native SKALA.
Definition skala_gpw_functional.F:144
skala_gpw_functional::get_gauxc_section
type(section_vals_type) function, pointer, public get_gauxc_section(xc_section)
Return the first GAUXC functional subsection, if present.
Definition skala_gpw_functional.F:2176
skala_gpw_functional::skala_gapw_atom_vxc_of_r
subroutine, public skala_gapw_atom_vxc_of_r(xc_section, grid_atom, group, atom_coord, rho, drho, tau, weights, lsd, nspins, na, nr, exc, vxc, vxg, vtau, energy_only, atom_force, atom_virial)
Evaluate SKALA on a GAPW one-center atomic grid.
Definition skala_gpw_functional.F:653
skala_gpw_functional::skala_gpw_exc_density
subroutine, public skala_gpw_exc_density(exc_r, rho_r, rho_g, tau, xc_section, weights, pw_pool, particle_set, cell)
Evaluate the native SKALA XC energy density on the CP2K PW grid.
Definition skala_gpw_functional.F:513
skala_gpw_functional::xc_section_uses_native_skala_grid
logical function, public xc_section_uses_native_skala_grid(xc_section)
Return true if the GAUXC subsection requests the CP2K-native GPW grid path.
Definition skala_gpw_functional.F:98
skala_gpw_functional::skala_gapw_density_partition_hard_only
integer, parameter, public skala_gapw_density_partition_hard_only
Definition skala_gpw_functional.F:72
skala_torch_api
Small CP2K wrapper around the SKALA TorchScript functional protocol.
Definition skala_torch_api.F:11
skala_torch_api::skala_torch_model_release
subroutine, public skala_torch_model_release(model)
Release a loaded SKALA TorchScript model.
Definition skala_torch_api.F:77
skala_torch_api::skala_torch_model_get_exc
subroutine, public skala_torch_model_get_exc(model, inputs, grid_weights, exc_tensor, exc)
Evaluate the weighted SKALA exchange-correlation energy.
Definition skala_torch_api.F:162
skala_torch_api::skala_torch_model_get_exc_density
subroutine, public skala_torch_model_get_exc_density(model, inputs, exc_density)
Evaluate the SKALA exchange-correlation energy density.
Definition skala_torch_api.F:136
skala_torch_api::skala_torch_model_load
subroutine, public skala_torch_model_load(model, filename)
Load a SKALA TorchScript model and its feature metadata.
Definition skala_torch_api.F:53
string_utilities
Utilities for string manipulations.
Definition string_utilities.F:16
string_utilities::uppercase
elemental subroutine, public uppercase(string)
Convert all lower case characters in a string to upper case.
Definition string_utilities.F:3362
torch_api
Definition torch_api.F:7
torch_api::torch_dict_release
subroutine, public torch_dict_release(dict)
Releases a Torch dictionary and all its ressources.
Definition torch_api.F:1799
torch_api::torch_use_cuda
subroutine, public torch_use_cuda(use_cuda)
Select whether Torch wrappers should use CUDA when available.
Definition torch_api.F:1551
torch_api::torch_tensor_backward_scalar
subroutine, public torch_tensor_backward_scalar(tensor)
Runs autograd on a scalar Torch tensor.
Definition torch_api.F:1500
torch_api::torch_tensor_to_device_leaf
subroutine, public torch_tensor_to_device_leaf(tensor, requires_grad)
Moves a tensor to the active Torch device and makes it an autograd leaf.
Definition torch_api.F:1523
torch_api::torch_dict_create
subroutine, public torch_dict_create(dict)
Creates an empty Torch dictionary.
Definition torch_api.F:1682
torch_api::torch_tensor_grad
subroutine, public torch_tensor_grad(tensor, grad)
Returns the gradient of a Torch tensor which was computed by autograd.
Definition torch_api.F:1572
torch_api::torch_cuda_device_count
integer function, public torch_cuda_device_count()
Return the number of CUDA devices visible to Torch.
Definition torch_api.F:2065
torch_api::torch_dict_insert
subroutine, public torch_dict_insert(dict, key, tensor)
Inserts a Torch tensor into a Torch dictionary.
Definition torch_api.F:1733
torch_api::torch_cuda_is_available
logical function, public torch_cuda_is_available()
Returns true iff the Torch CUDA backend is available.
Definition torch_api.F:2044
torch_api::torch_tensor_release
subroutine, public torch_tensor_release(tensor)
Releases a Torch tensor and all its ressources.
Definition torch_api.F:1658
xc_rho_cflags_types
contains the structure
Definition xc_rho_cflags_types.F:14
xc_rho_set_types
contains the structure
Definition xc_rho_set_types.F:14
xc_rho_set_types::xc_rho_set_create
subroutine, public xc_rho_set_create(rho_set, local_bounds, rho_cutoff, drho_cutoff, tau_cutoff)
allocates and does (minimal) initialization of a rho_set
Definition xc_rho_set_types.F:111
xc_rho_set_types::xc_rho_set_release
subroutine, public xc_rho_set_release(rho_set, pw_pool)
releases the given rho_set
Definition xc_rho_set_types.F:129
xc_rho_set_types::xc_rho_set_update
subroutine, public xc_rho_set_update(rho_set, rho_r, rho_g, tau, needs, xc_deriv_method_id, xc_rho_smooth_id, pw_pool, spinflip)
updates the given rho set with the density given by rho_r (and rho_g). The rho set will contain the c...
Definition xc_rho_set_types.F:691
xc_rho_set_types::xc_rho_set_get
subroutine, public xc_rho_set_get(rho_set, can_return_null, rho, drho, norm_drho, rhoa, rhob, norm_drhoa, norm_drhob, rho_1_3, rhoa_1_3, rhob_1_3, laplace_rho, laplace_rhoa, laplace_rhob, drhoa, drhob, rho_cutoff, drho_cutoff, tau_cutoff, tau, tau_a, tau_b, local_bounds)
returns the various attributes of rho_set
Definition xc_rho_set_types.F:281
xc_util
contains utility functions for the xc package
Definition xc_util.F:14
xc_util::xc_pw_divergence
subroutine, public xc_pw_divergence(xc_deriv_method_id, pw_to_deriv, tmp_g, vxc_g, vxc_r)
Calculates the divergence of pw_to_deriv.
Definition xc_util.F:253
xc_util::xc_requires_tmp_g
elemental logical function, public xc_requires_tmp_g(xc_deriv_id)
...
Definition xc_util.F:58
cell_types::cell_type
Type defining parameters related to the simulation cell.
Definition cell_types.F:60
cp_array_utils::cp_3d_r_cp_type
represent a pointer to a contiguous 3d array
Definition cp_array_utils.F:149
input_section_types::section_vals_type
stores the values of a section
Definition input_section_types.F:127
message_passing::mp_comm_type
Definition message_passing.F:158
particle_types::particle_type
Definition particle_types.F:35
pw_grid_types::pw_grid_type
Definition pw_grid_types.F:53
pw_pool_types::pw_pool_type
Manages a pool of grids (to be used for example as tmp objects), but can also be used to instantiate ...
Definition pw_pool_types.F:83
pw_types::pw_c1d_gs_type
Definition pw_types.F:127
pw_types::pw_r3d_rs_type
Definition pw_types.F:62
qs_grid_atom::grid_atom_type
Definition qs_grid_atom.F:44
skala_gpw_features::skala_gpw_feature_type
Definition skala_gpw_features.F:124
skala_torch_api::skala_torch_model_type
Definition skala_torch_api.F:37
torch_api::torch_dict_type
Definition torch_api.F:34
torch_api::torch_tensor_type
Definition torch_api.F:29
xc_rho_cflags_types::xc_rho_cflags_type
contains a flag for each component of xc_rho_set, so that you can use it to tell which components you...
Definition xc_rho_cflags_types.F:48
xc_rho_set_types::xc_rho_set_type
represent a density, with all the representation and data needed to perform a functional evaluation
Definition xc_rho_set_types.F:78