(git:6344619)
Loading...
Searching...
No Matches
skala_gpw_features.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 Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
10! **************************************************************************************************
11MODULE skala_gpw_features
12 USE cell_types, ONLY: cell_type,&
13 pbc
14 USE cp_array_utils, ONLY: cp_3d_r_cp_type
15 USE kinds, ONLY: dp,&
16 int_8
17 USE message_passing, ONLY: mp_comm_type
18 USE particle_types, ONLY: particle_type
19 USE pw_grid_types, ONLY: pw_grid_type
20 USE pw_types, ONLY: pw_r3d_rs_type
21 USE torch_api, ONLY: &
22 torch_dict_clone, torch_dict_create, torch_dict_insert, torch_dict_release, &
23 torch_dict_type, torch_tensor_expand_dim, torch_tensor_from_array, torch_tensor_narrow, &
24 torch_tensor_release, torch_tensor_reset_from_array, torch_tensor_to_device_leaf, &
25 torch_tensor_type
26 USE xc_rho_set_types, ONLY: xc_rho_set_get,&
27 xc_rho_set_type
28#include "../base/base_uses.f90"
29
30 IMPLICIT NONE
31
32 PRIVATE
33
34 CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'skala_gpw_features'
35 REAL(KIND=dp), PARAMETER, PRIVATE :: layout_tol = 1.0e-12_dp
36 INTEGER, PARAMETER, PRIVATE :: ndynamic_per_point = 10, nrks_dynamic_per_point = 5, &
37 nstatic_per_point = 5, ngrad_per_point = 10
38 INTEGER, PARAMETER, PUBLIC :: skala_gpw_atom_partition_hard = 1, &
39 skala_gpw_atom_partition_smooth = 2
40 REAL(kind=dp), PARAMETER, PRIVATE :: smooth_partition_eps = 1.0e-12_dp
41 REAL(kind=dp), PARAMETER, PRIVATE :: smooth_partition_taper_end = &
42 10.0_dp*smooth_partition_eps
43 INTEGER, PARAMETER, PRIVATE :: periodic_partition_image_shell = 1
44
45 PUBLIC :: skala_gpw_atom_subchunk_count, skala_gpw_atom_subchunk_layout, &
46 skala_gpw_feature_build, skala_gpw_feature_build_atom_subchunk, &
47 skala_gpw_feature_build_atom_subchunk_bounds, skala_gpw_feature_release, &
48 skala_gpw_feature_type, skala_gpw_smooth_partition_derivatives, &
49 build_periodic_atom_image_layout, periodic_atom_image_partition, &
50 periodic_atom_image_partition_from_layout, smooth_atom_partition, &
51 smooth_partition_atomic_weight_scale, smooth_partition_atomic_weight_scale_derivative
52
53 TYPE skala_gpw_layout_cache_type
54 INTEGER :: chunk_atom_begin = 1, chunk_atom_end = 0, &
55 chunk_feature_begin = 1, &
56 chunk_feature_count = 0, chunk_natom = 0, &
57 natom = 0, nflat = 0, nflat_local = 0, &
58 npoint = 0, nproc = 0, &
59 atom_partition = skala_gpw_atom_partition_hard
60 INTEGER, DIMENSION(2, 3) :: bo = 0, bounds = 0
61 INTEGER, DIMENSION(3) :: npts = 0
62 INTEGER, ALLOCATABLE, DIMENSION(:) :: dynamic_counts, dynamic_displs, &
63 chunk_feature_counts, chunk_feature_displs, &
64 chunk_grad_counts, chunk_grad_displs, &
65 feature_counts, feature_displs, &
66 feature_source_points, global_to_feature, &
67 local_feature_counts, local_feature_offsets, &
68 local_feature_points, local_feature_rows, &
69 chunk_return_positions, route_chunk_offsets, &
70 route_chunk_rows, route_row_offsets, &
71 route_row_positions, route_send_local_rows, &
72 route_point_recv_counts, &
73 route_point_recv_displs, &
74 route_point_send_counts, &
75 route_point_send_displs
76 INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
77 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes, chunk_atomic_grid_sizes, &
78 chunk_feature_indices
79 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: local_feature_indices
80 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape, &
81 chunk_atomic_grid_size_bound_shape
82 TYPE(torch_dict_type) :: chunk_inputs
83 TYPE(torch_dict_type) :: chunk_static_inputs
84 TYPE(torch_dict_type) :: inputs
85 TYPE(torch_dict_type) :: static_inputs
86 TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
87 TYPE(torch_tensor_type) :: atomic_grid_sizes_t
88 TYPE(torch_tensor_type) :: atomic_grid_weights_t
89 TYPE(torch_tensor_type) :: chunk_atomic_grid_size_bound_shape_t
90 TYPE(torch_tensor_type) :: chunk_atomic_grid_sizes_t
91 TYPE(torch_tensor_type) :: chunk_atomic_grid_weights_t
92 TYPE(torch_tensor_type) :: chunk_coarse_0_atomic_coords_t
93 TYPE(torch_tensor_type) :: chunk_density_t
94 TYPE(torch_tensor_type) :: chunk_density_input_t
95 TYPE(torch_tensor_type) :: chunk_feature_indices_t
96 TYPE(torch_tensor_type) :: chunk_grad_t
97 TYPE(torch_tensor_type) :: chunk_grad_input_t
98 TYPE(torch_tensor_type) :: chunk_grid_coords_t
99 TYPE(torch_tensor_type) :: chunk_grid_weights_t
100 TYPE(torch_tensor_type) :: chunk_kin_t
101 TYPE(torch_tensor_type) :: chunk_kin_input_t
102 TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
103 TYPE(torch_tensor_type) :: density_t
104 TYPE(torch_tensor_type) :: grid_coords_t
105 TYPE(torch_tensor_type) :: grid_weights_t
106 TYPE(torch_tensor_type) :: grad_t
107 TYPE(torch_tensor_type) :: kin_t
108 TYPE(torch_tensor_type) :: local_feature_indices_t
109 REAL(kind=dp) :: dvol = 0.0_dp, weight_sum = 0.0_dp, &
110 weight_sumsq = 0.0_dp
111 REAL(kind=dp), DIMENSION(3, 3) :: cell_hmat = 0.0_dp, dh = 0.0_dp
112 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, chunk_atomic_grid_weights, &
113 chunk_grid_weights, grid_weights
114 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords, chunk_coarse_0_atomic_coords, &
115 chunk_grid_coords, coarse_0_atomic_coords, &
116 grid_coords
117 LOGICAL :: active = .false., has_weights = .false., &
118 chunk_dynamic_input_views_active = .false., &
119 chunk_dynamic_tensors_active = .false., &
120 chunk_inputs_active = .false., &
121 chunk_inputs_use_collapsed_rks = .false., &
122 chunk_static_tensors_active = .false., &
123 dynamic_tensors_active = .false., &
124 inputs_active = .false., &
125 static_tensors_active = .false.
126 END TYPE skala_gpw_layout_cache_type
127
128 TYPE skala_gpw_feature_type
129 INTEGER :: chunk_feature_count = 0, nflat = 0, &
130 nflat_local = 0, &
131 atom_partition = skala_gpw_atom_partition_hard
132 TYPE(torch_dict_type) :: inputs
133 TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
134 TYPE(torch_tensor_type) :: atomic_grid_sizes_t
135 TYPE(torch_tensor_type) :: atomic_grid_weights_t
136 TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
137 TYPE(torch_tensor_type) :: density_input_t
138 TYPE(torch_tensor_type) :: density_t
139 TYPE(torch_tensor_type) :: grad_t
140 TYPE(torch_tensor_type) :: grad_input_t
141 TYPE(torch_tensor_type) :: grid_coords_t
142 TYPE(torch_tensor_type) :: grid_weights_t
143 TYPE(torch_tensor_type) :: kin_input_t
144 TYPE(torch_tensor_type) :: kin_t
145 TYPE(torch_tensor_type) :: local_feature_indices_t
146 INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_grad_counts, chunk_grad_displs, &
147 local_feature_counts, local_feature_offsets, &
148 local_feature_rows
149 INTEGER, POINTER, DIMENSION(:) :: route_point_recv_counts => null(), &
150 route_point_recv_displs => null(), &
151 route_point_send_counts => null(), &
152 route_point_send_displs => null(), &
153 route_chunk_offsets => null(), &
154 route_chunk_rows => null(), &
155 route_row_offsets => null(), &
156 route_row_positions => null(), &
157 route_send_local_rows => null()
158 INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
159 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
160 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
161 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
162 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: chunk_density, chunk_kin, &
163 coarse_0_atomic_coords, density, &
164 grid_coords, kin
165 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: chunk_grad, grad
166 REAL(kind=dp) :: electron_count = 0.0_dp, &
167 grid_weight_sum = 0.0_dp, &
168 kinetic_integral = 0.0_dp, &
169 spin_moment = 0.0_dp
170 LOGICAL :: active = .false., owns_coordinate_tensor = .false., &
171 owns_grid_coordinate_tensor = .false., &
172 owns_weight_tensors = .false., &
173 owns_dynamic_tensors = .true., &
174 owns_inputs = .true., &
175 owns_static_tensors = .true., &
176 uses_atom_chunk_routing = .false., &
177 uses_atom_chunks = .false., &
178 uses_collapsed_rks_dynamic = .false.
179 END TYPE skala_gpw_feature_type
180
181 TYPE(skala_gpw_layout_cache_type), SAVE, TARGET :: cached_layout
182
183CONTAINS
184
185! **************************************************************************************************
186!> \brief Build a flat SKALA molecular feature dictionary from a local GPW grid.
187!> \param features ...
188!> \param rho_set ...
189!> \param rho_r ...
190!> \param particle_set ...
191!> \param cell ...
192!> \param requires_grad ...
193!> \param weights ...
194!> \param requires_coordinate_grad ...
195!> \param requires_stress_grad ...
196!> \param use_atom_chunks ...
197!> \param route_atom_chunks ...
198!> \param atom_partition ...
199! **************************************************************************************************
200 SUBROUTINE skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
201 requires_grad, weights, requires_coordinate_grad, &
202 requires_stress_grad, use_atom_chunks, route_atom_chunks, &
203 atom_partition)
204 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
205 TYPE(xc_rho_set_type), INTENT(IN) :: rho_set
206 TYPE(pw_r3d_rs_type), DIMENSION(:), INTENT(IN) :: rho_r
207 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
208 TYPE(cell_type), POINTER :: cell
209 LOGICAL, INTENT(IN), OPTIONAL :: requires_grad
210 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
211 LOGICAL, INTENT(IN), OPTIONAL :: requires_coordinate_grad, &
212 requires_stress_grad, use_atom_chunks, &
213 route_atom_chunks
214 INTEGER, INTENT(IN), OPTIONAL :: atom_partition
215
216 INTEGER :: handle, i, ipt, ispin, j, k, local_row, my_atom_partition, &
217 ndynamic_local_per_point, nflat, nflat_local, nspins, phase_handle, real_base, row
218 INTEGER, DIMENSION(2, 3) :: bo
219 LOGICAL :: collapse_spin_dynamics, my_requires_coordinate_grad, my_requires_grad, &
220 my_requires_stress_grad, my_route_atom_chunks, my_use_atom_chunks, &
221 use_atom_chunk_protocol, use_atom_chunk_routing
222 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: global_dynamic, local_dynamic
223 REAL(kind=dp), DIMENSION(3) :: integrals
224 REAL(kind=dp), DIMENSION(:, :, :), POINTER :: rho, rhoa, rhob, tau_a, tau_b, tau_total
225 TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
226 TYPE(pw_grid_type), POINTER :: pw_grid
227
228 CALL timeset("skala_gpw_feature_build", handle)
229
230 my_requires_grad = .false.
231 IF (PRESENT(requires_grad)) my_requires_grad = requires_grad
232 my_requires_coordinate_grad = .false.
233 IF (PRESENT(requires_coordinate_grad)) THEN
234 my_requires_coordinate_grad = requires_coordinate_grad
235 END IF
236 my_requires_stress_grad = .false.
237 IF (PRESENT(requires_stress_grad)) my_requires_stress_grad = requires_stress_grad
238 my_use_atom_chunks = .false.
239 IF (PRESENT(use_atom_chunks)) my_use_atom_chunks = use_atom_chunks
240 my_route_atom_chunks = .false.
241 IF (PRESENT(route_atom_chunks)) my_route_atom_chunks = route_atom_chunks
242 my_atom_partition = skala_gpw_atom_partition_hard
243 IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
244 IF (my_atom_partition /= skala_gpw_atom_partition_hard .AND. &
245 my_atom_partition /= skala_gpw_atom_partition_smooth) THEN
246 CALL cp_abort(__location__, "Unknown native SKALA atom-partition mode.")
247 END IF
248 cpassert(ASSOCIATED(cell))
249 cpassert(ASSOCIATED(particle_set))
250 cpassert(SIZE(rho_r) == 1 .OR. SIZE(rho_r) == 2)
251 cpassert(ASSOCIATED(rho_r(1)%pw_grid))
252 pw_grid => rho_r(1)%pw_grid
253
254 nspins = SIZE(rho_r)
255 bo = pw_grid%bounds_local
256 nflat_local = pw_grid%ngpts_local
257
258 CALL timeset("skala_gpw_pre_release", phase_handle)
259 CALL skala_gpw_feature_release(features)
260 CALL timestop(phase_handle)
261
262 use_atom_chunk_protocol = my_use_atom_chunks .AND. &
263 .NOT. (my_requires_coordinate_grad .OR. my_requires_stress_grad)
264 use_atom_chunk_routing = use_atom_chunk_protocol .AND. my_route_atom_chunks
265 CALL timeset("skala_gpw_layout_cache", phase_handle)
266 CALL ensure_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition, &
267 needs_full_static_tensors=.NOT. use_atom_chunk_protocol)
268 CALL timestop(phase_handle)
269 nflat = cached_layout%nflat
270 collapse_spin_dynamics = nspins == 1 .AND. use_atom_chunk_routing
271 ndynamic_local_per_point = ndynamic_per_point
272 IF (collapse_spin_dynamics) ndynamic_local_per_point = nrks_dynamic_per_point
273 ALLOCATE (local_dynamic(ndynamic_local_per_point*nflat_local))
274 local_dynamic = 0.0_dp
275
276 CALL timeset("skala_gpw_pack_local", phase_handle)
277 IF (nspins == 1) THEN
278 CALL xc_rho_set_get(rho_set, rho=rho, drho=drho, tau=tau_total)
279 ELSE
280 CALL xc_rho_set_get(rho_set, rhoa=rhoa, rhob=rhob, drhoa=drhoa, drhob=drhob, &
281 tau_a=tau_a, tau_b=tau_b)
282 END IF
283
284 local_row = 0
285 DO k = bo(1, 3), bo(2, 3)
286 DO j = bo(1, 2), bo(2, 2)
287 DO i = bo(1, 1), bo(2, 1)
288 local_row = local_row + 1
289 real_base = ndynamic_local_per_point*(local_row - 1)
290
291 IF (nspins == 1) THEN
292 IF (collapse_spin_dynamics) THEN
293 local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
294 local_dynamic(real_base + 2) = 0.5_dp*drho(1)%array(i, j, k)
295 local_dynamic(real_base + 3) = 0.5_dp*drho(2)%array(i, j, k)
296 local_dynamic(real_base + 4) = 0.5_dp*drho(3)%array(i, j, k)
297 local_dynamic(real_base + 5) = 0.5_dp*tau_total(i, j, k)
298 ELSE
299 local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
300 local_dynamic(real_base + 2) = 0.5_dp*rho(i, j, k)
301 DO ispin = 1, 2
302 local_dynamic(real_base + 2 + 3*(ispin - 1) + 1) = &
303 0.5_dp*drho(1)%array(i, j, k)
304 local_dynamic(real_base + 2 + 3*(ispin - 1) + 2) = &
305 0.5_dp*drho(2)%array(i, j, k)
306 local_dynamic(real_base + 2 + 3*(ispin - 1) + 3) = &
307 0.5_dp*drho(3)%array(i, j, k)
308 local_dynamic(real_base + 8 + ispin) = 0.5_dp*tau_total(i, j, k)
309 END DO
310 END IF
311 ELSE
312 local_dynamic(real_base + 1) = rhoa(i, j, k)
313 local_dynamic(real_base + 2) = rhob(i, j, k)
314 local_dynamic(real_base + 3) = drhoa(1)%array(i, j, k)
315 local_dynamic(real_base + 4) = drhoa(2)%array(i, j, k)
316 local_dynamic(real_base + 5) = drhoa(3)%array(i, j, k)
317 local_dynamic(real_base + 6) = drhob(1)%array(i, j, k)
318 local_dynamic(real_base + 7) = drhob(2)%array(i, j, k)
319 local_dynamic(real_base + 8) = drhob(3)%array(i, j, k)
320 local_dynamic(real_base + 9) = tau_a(i, j, k)
321 local_dynamic(real_base + 10) = tau_b(i, j, k)
322 END IF
323 END DO
324 END DO
325 END DO
326 CALL timestop(phase_handle)
327
328 CALL timeset("skala_gpw_copy_layout", phase_handle)
329 CALL copy_cached_layout(features, use_atom_chunk_routing, &
330 my_requires_coordinate_grad .OR. my_requires_stress_grad, &
331 my_requires_stress_grad .OR. &
332 (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
333 (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
334 CALL timestop(phase_handle)
335
336 IF (use_atom_chunk_routing) THEN
337 CALL timeset("skala_gpw_route_dyn", phase_handle)
338 CALL route_atom_chunk_dynamics(features, local_dynamic, pw_grid%para%group, &
339 collapse_spin_dynamics)
340 features%uses_atom_chunk_routing = .true.
341 features%uses_atom_chunks = .true.
342 CALL timestop(phase_handle)
343 ELSE
344 ALLOCATE (global_dynamic(ndynamic_per_point*cached_layout%npoint))
345 CALL timeset("skala_gpw_allgatherv", phase_handle)
346 CALL pw_grid%para%group%allgatherv(local_dynamic, global_dynamic, &
347 cached_layout%dynamic_counts, &
348 cached_layout%dynamic_displs)
349 CALL timestop(phase_handle)
350
351 CALL timeset("skala_gpw_reorder_dyn", phase_handle)
352 ALLOCATE (features%density(nflat, 2), features%grad(nflat, 3, 2), &
353 features%kin(nflat, 2))
354 features%density = 0.0_dp
355 features%grad = 0.0_dp
356 features%kin = 0.0_dp
357
358 DO row = 1, nflat
359 ipt = cached_layout%feature_source_points(row)
360 real_base = ndynamic_per_point*(ipt - 1)
361 features%density(row, :) = global_dynamic(real_base + 1:real_base + 2)
362 features%grad(row, 1, 1) = global_dynamic(real_base + 3)
363 features%grad(row, 2, 1) = global_dynamic(real_base + 4)
364 features%grad(row, 3, 1) = global_dynamic(real_base + 5)
365 features%grad(row, 1, 2) = global_dynamic(real_base + 6)
366 features%grad(row, 2, 2) = global_dynamic(real_base + 7)
367 features%grad(row, 3, 2) = global_dynamic(real_base + 8)
368 features%kin(row, :) = global_dynamic(real_base + 9:real_base + 10)
369 END DO
370 CALL timestop(phase_handle)
371 END IF
372
373 CALL timeset("skala_gpw_feature_sums", phase_handle)
374 IF (features%uses_atom_chunks) THEN
375 features%electron_count = 0.0_dp
376 features%kinetic_integral = 0.0_dp
377 features%spin_moment = 0.0_dp
378 IF (features%chunk_feature_count > 0) THEN
379 IF (features%uses_collapsed_rks_dynamic) THEN
380 features%electron_count = sum(2.0_dp*features%chunk_density(:, 1)* &
381 cached_layout%chunk_grid_weights)
382 features%kinetic_integral = sum(2.0_dp*features%chunk_kin(:, 1)* &
383 cached_layout%chunk_grid_weights)
384 ELSE
385 features%electron_count = sum((features%chunk_density(:, 1) + &
386 features%chunk_density(:, 2))* &
387 cached_layout%chunk_grid_weights)
388 features%kinetic_integral = sum((features%chunk_kin(:, 1) + &
389 features%chunk_kin(:, 2))* &
390 cached_layout%chunk_grid_weights)
391 features%spin_moment = sum((features%chunk_density(:, 1) - &
392 features%chunk_density(:, 2))* &
393 cached_layout%chunk_grid_weights)
394 END IF
395 END IF
396 integrals = [features%electron_count, features%kinetic_integral, features%spin_moment]
397 CALL pw_grid%para%group%sum(integrals)
398 features%electron_count = integrals(1)
399 features%kinetic_integral = integrals(2)
400 features%spin_moment = integrals(3)
401 ELSE
402 features%electron_count = sum((features%density(:, 1) + features%density(:, 2))* &
403 features%grid_weights)
404 features%kinetic_integral = sum((features%kin(:, 1) + features%kin(:, 2))* &
405 features%grid_weights)
406 features%spin_moment = sum((features%density(:, 1) - features%density(:, 2))* &
407 features%grid_weights)
408 END IF
409 IF (use_atom_chunk_routing) THEN
410 features%grid_weight_sum = sum(cached_layout%grid_weights)
411 ELSE
412 features%grid_weight_sum = sum(features%grid_weights)
413 END IF
414 CALL timestop(phase_handle)
415
416 CALL timeset("skala_gpw_tensor_update", phase_handle)
417 IF (use_atom_chunk_protocol .AND. .NOT. features%uses_atom_chunks) THEN
418 IF (features%chunk_feature_count > 0) CALL extract_atom_chunk_dynamics(features)
419 features%uses_atom_chunks = .true.
420 END IF
421 IF (.NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0) THEN
422 CALL add_feature_tensors(features, my_requires_grad, my_requires_coordinate_grad, &
423 my_requires_stress_grad, &
424 features%uses_atom_chunks, &
425 requires_weight_grad= &
426 (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
427 (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
428 ELSE
429 ! This rank participates in atom-chunk communication but owns no model input rows.
430 features%owns_coordinate_tensor = .false.
431 features%owns_grid_coordinate_tensor = .false.
432 features%owns_weight_tensors = .false.
433 features%owns_dynamic_tensors = .false.
434 features%owns_inputs = .false.
435 features%owns_static_tensors = .false.
436 END IF
437 CALL timestop(phase_handle)
438 features%active = .true.
439
440 IF (ALLOCATED(global_dynamic)) DEALLOCATE (global_dynamic)
441 DEALLOCATE (local_dynamic)
442 CALL timestop(handle)
443
444 END SUBROUTINE skala_gpw_feature_build
445
446! **************************************************************************************************
447!> \brief Ensure that static grid-to-atom layout data is cached for the current grid/geometry.
448!> \param pw_grid ...
449!> \param particle_set ...
450!> \param cell ...
451!> \param weights ...
452!> \param atom_partition ...
453!> \param needs_full_static_tensors ...
454! **************************************************************************************************
455 SUBROUTINE ensure_layout_cache(pw_grid, particle_set, cell, weights, atom_partition, &
456 needs_full_static_tensors)
457 TYPE(pw_grid_type), POINTER :: pw_grid
458 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
459 TYPE(cell_type), POINTER :: cell
460 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
461 INTEGER, INTENT(IN), OPTIONAL :: atom_partition
462 LOGICAL, INTENT(IN) :: needs_full_static_tensors
463
464 INTEGER :: my_atom_partition, phase_handle
465 LOGICAL :: cache_matches
466
467 my_atom_partition = skala_gpw_atom_partition_hard
468 IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
469 IF (PRESENT(weights)) THEN
470 CALL timeset("skala_gpw_layout_match", phase_handle)
471 cache_matches = layout_cache_matches(pw_grid, particle_set, cell, weights, &
472 my_atom_partition)
473 CALL timestop(phase_handle)
474 ELSE
475 CALL timeset("skala_gpw_layout_match", phase_handle)
476 cache_matches = layout_cache_matches(pw_grid, particle_set, cell, &
477 atom_partition=my_atom_partition)
478 CALL timestop(phase_handle)
479 END IF
480
481 IF (cache_matches) THEN
482 IF (needs_full_static_tensors .AND. .NOT. cached_layout%static_tensors_active) THEN
483 CALL timeset("skala_gpw_layout_tensors", phase_handle)
484 CALL build_full_static_layout_tensors(cached_layout)
485 CALL timestop(phase_handle)
486 END IF
487 RETURN
488 END IF
489
490 CALL timeset("skala_gpw_layout_rebuild", phase_handle)
491 IF (PRESENT(weights)) THEN
492 CALL rebuild_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition, &
493 needs_full_static_tensors)
494 ELSE
495 CALL rebuild_layout_cache(pw_grid, particle_set, cell, &
496 atom_partition=my_atom_partition, &
497 needs_full_static_tensors=needs_full_static_tensors)
498 END IF
499 CALL timestop(phase_handle)
500
501 END SUBROUTINE ensure_layout_cache
502
503! **************************************************************************************************
504!> \brief Check whether the current static layout cache can be reused.
505!> \param pw_grid ...
506!> \param particle_set ...
507!> \param cell ...
508!> \param weights ...
509!> \param atom_partition ...
510!> \return ...
511! **************************************************************************************************
512 FUNCTION layout_cache_matches(pw_grid, particle_set, cell, weights, atom_partition) RESULT(matches)
513 TYPE(pw_grid_type), POINTER :: pw_grid
514 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
515 TYPE(cell_type), POINTER :: cell
516 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
517 INTEGER, INTENT(IN), OPTIONAL :: atom_partition
518 LOGICAL :: matches
519
520 INTEGER :: iatom, my_atom_partition
521 LOGICAL :: weights_match
522
523 my_atom_partition = skala_gpw_atom_partition_hard
524 IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
525 matches = .false.
526 IF (.NOT. cached_layout%active) RETURN
527 IF (cached_layout%atom_partition /= my_atom_partition) RETURN
528 IF (cached_layout%natom /= SIZE(particle_set)) RETURN
529 IF (cached_layout%nflat_local /= pw_grid%ngpts_local) RETURN
530 IF (cached_layout%nproc /= pw_grid%para%group%num_pe) RETURN
531 IF (any(cached_layout%bo /= pw_grid%bounds_local)) RETURN
532 IF (any(cached_layout%bounds /= pw_grid%bounds)) RETURN
533 IF (any(cached_layout%npts /= pw_grid%npts)) RETURN
534 IF (abs(cached_layout%dvol - pw_grid%dvol) > layout_tol) RETURN
535 IF (any(abs(cached_layout%dh - pw_grid%dh) > layout_tol)) RETURN
536 IF (any(abs(cached_layout%cell_hmat - cell%hmat) > layout_tol)) RETURN
537 IF (.NOT. ALLOCATED(cached_layout%atom_coords)) RETURN
538
539 DO iatom = 1, SIZE(particle_set)
540 IF (any(abs(cached_layout%atom_coords(:, iatom) - particle_set(iatom)%r) > layout_tol)) RETURN
541 END DO
542
543 IF (PRESENT(weights)) THEN
544 weights_match = layout_weights_match(pw_grid, weights)
545 ELSE
546 weights_match = layout_weights_match(pw_grid)
547 END IF
548 IF (.NOT. weights_match) RETURN
549
550 matches = .true.
551
552 END FUNCTION layout_cache_matches
553
554! **************************************************************************************************
555!> \brief Check whether current optional integration weights match the cached static tensors.
556!> \param pw_grid ...
557!> \param weights ...
558!> \return ...
559! **************************************************************************************************
560 FUNCTION layout_weights_match(pw_grid, weights) RESULT(matches)
561 TYPE(pw_grid_type), POINTER :: pw_grid
562 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
563 LOGICAL :: matches
564
565 LOGICAL :: has_weights
566 REAL(kind=dp) :: weight_sum, weight_sumsq
567
568 matches = .false.
569 mark_used(pw_grid)
570 IF (PRESENT(weights)) THEN
571 CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
572 ELSE
573 CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
574 weight_sumsq=weight_sumsq)
575 END IF
576
577 IF (cached_layout%has_weights .NEQV. has_weights) RETURN
578 IF (abs(cached_layout%weight_sum - weight_sum) > layout_tol) RETURN
579 IF (abs(cached_layout%weight_sumsq - weight_sumsq) > layout_tol) RETURN
580
581 matches = .true.
582
583 END FUNCTION layout_weights_match
584
585! **************************************************************************************************
586!> \brief Build the static SKALA layout cache.
587!> \param pw_grid ...
588!> \param particle_set ...
589!> \param cell ...
590!> \param weights ...
591!> \param atom_partition ...
592!> \param needs_full_static_tensors ...
593! **************************************************************************************************
594 SUBROUTINE rebuild_layout_cache(pw_grid, particle_set, cell, weights, atom_partition, &
595 needs_full_static_tensors)
596 TYPE(pw_grid_type), POINTER :: pw_grid
597 TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
598 TYPE(cell_type), POINTER :: cell
599 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
600 INTEGER, INTENT(IN), OPTIONAL :: atom_partition
601 LOGICAL, INTENT(IN) :: needs_full_static_tensors
602
603 INTEGER :: feature_local, feature_slot, i, iatom, ipt, j, jatom, k, local_feature, &
604 local_row, max_grid_size, max_local_features, my_atom_partition, natom, nfeature_local, &
605 nflat, nflat_local, npoint, nproc, nx, ny, owner, pe, pe_index, phase_handle, row, &
606 source_global, source_local, static_base
607 INTEGER, ALLOCATABLE, DIMENSION(:) :: atom_offset, atom_position, chunk_atom_begin, &
608 chunk_atom_end, cursor, feature_counts, feature_displs, global_owner, &
609 global_source_points, local_feature_counts_tmp, local_owner, local_source_global, &
610 local_source_points, point_counts, point_displs, static_counts, static_displs
611 INTEGER, DIMENSION(2, 3) :: bo
612 LOGICAL :: has_weights, nonperiodic
613 REAL(kind=dp) :: atomic_base_weight, base_weight, &
614 included_sum, partition_weight, &
615 weight_sum, weight_sumsq
616 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: distances, global_static, local_static, &
617 partition_weights
618 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc, pair_distances, &
619 partition_atom_coords
620 REAL(kind=dp), DIMENSION(3) :: grid_point
621
622 CALL release_layout_cache(cached_layout)
623
624 my_atom_partition = skala_gpw_atom_partition_hard
625 IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
626 natom = SIZE(particle_set)
627 bo = pw_grid%bounds_local
628 nflat_local = pw_grid%ngpts_local
629 nproc = pw_grid%para%group%num_pe
630 pe_index = pw_grid%para%group%mepos + 1
631
632 IF (PRESENT(weights)) THEN
633 CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
634 ELSE
635 CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
636 weight_sumsq=weight_sumsq)
637 END IF
638
639 max_local_features = nflat_local
640 IF (my_atom_partition == skala_gpw_atom_partition_smooth) THEN
641 max_local_features = nflat_local*natom
642 END IF
643 ALLOCATE (local_owner(max_local_features), &
644 local_source_points(max_local_features), &
645 local_static(nstatic_per_point*max_local_features), &
646 local_feature_counts_tmp(nflat_local), feature_counts(nproc), &
647 feature_displs(nproc), point_counts(nproc), point_displs(nproc), &
648 static_counts(nproc), static_displs(nproc), atom_coords_pbc(3, natom))
649 ALLOCATE (cached_layout%feature_index(bo(1, 1):bo(2, 1), &
650 bo(1, 2):bo(2, 2), &
651 bo(1, 3):bo(2, 3)))
652 cached_layout%feature_index = 0
653 local_feature_counts_tmp = 0
654 DO iatom = 1, natom
655 atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.true.)
656 END DO
657 nonperiodic = all(cell%perd == 0)
658 IF (my_atom_partition == skala_gpw_atom_partition_smooth .AND. nonperiodic) THEN
659 ALLOCATE (pair_distances(natom, natom))
660 pair_distances = 0.0_dp
661 DO iatom = 1, natom - 1
662 DO jatom = iatom + 1, natom
663 pair_distances(iatom, jatom) = &
664 sqrt(sum((atom_coords_pbc(:, iatom) - atom_coords_pbc(:, jatom))**2))
665 pair_distances(jatom, iatom) = pair_distances(iatom, jatom)
666 END DO
667 END DO
668 ELSE
669 ALLOCATE (pair_distances(0, 0))
670 END IF
671
672 CALL timeset("skala_gpw_layout_local", phase_handle)
673 nx = bo(2, 1) - bo(1, 1) + 1
674 ny = bo(2, 2) - bo(1, 2) + 1
675 IF (my_atom_partition == skala_gpw_atom_partition_hard) THEN
676 nfeature_local = nflat_local
677!$OMP PARALLEL DO DEFAULT(NONE) COLLAPSE(3) &
678!$OMP SHARED(atom_coords_pbc, bo, cached_layout, cell, local_feature_counts_tmp, local_owner, &
679!$OMP local_source_points, local_static, nx, ny, pw_grid, weights) &
680!$OMP PRIVATE(atomic_base_weight, base_weight, grid_point, i, j, k, local_row, owner, static_base)
681 DO k = bo(1, 3), bo(2, 3)
682 DO j = bo(1, 2), bo(2, 2)
683 DO i = bo(1, 1), bo(2, 1)
684 local_row = i - bo(1, 1) + 1 + &
685 nx*(j - bo(1, 2) + ny*(k - bo(1, 3)))
686 grid_point = grid_coordinate(pw_grid, [i, j, k])
687 atomic_base_weight = pw_grid%dvol
688 base_weight = pw_grid%dvol
689 IF (PRESENT(weights)) THEN
690 IF (ASSOCIATED(weights)) base_weight = base_weight*weights%array(i, j, k)
691 END IF
692 cached_layout%feature_index(i, j, k) = local_row
693 owner = nearest_atom(grid_point, atom_coords_pbc, cell)
694 local_feature_counts_tmp(local_row) = 1
695 local_owner(local_row) = owner
696 local_source_points(local_row) = local_row
697 static_base = nstatic_per_point*(local_row - 1)
698 local_static(static_base + 1:static_base + 3) = grid_point
699 local_static(static_base + 4) = base_weight
700 local_static(static_base + 5) = atomic_base_weight
701 END DO
702 END DO
703 END DO
704!$OMP END PARALLEL DO
705 ELSE
706!$OMP PARALLEL DEFAULT(NONE) &
707!$OMP SHARED(atom_coords_pbc, bo, cached_layout, cell, local_feature_counts_tmp, local_owner, &
708!$OMP local_source_points, local_static, natom, nonperiodic, nx, ny, pair_distances, &
709!$OMP pw_grid, weights) &
710!$OMP PRIVATE(atomic_base_weight, base_weight, distances, feature_slot, grid_point, i, iatom, &
711!$OMP included_sum, j, k, local_feature, local_row, owner, partition_weight, &
712!$OMP partition_atom_coords, partition_weights, static_base)
713 ! Keep point-local scratch outside the hot loop and private to each thread.
714 ALLOCATE (distances(natom), partition_atom_coords(3, natom), partition_weights(natom))
715!$OMP DO COLLAPSE(3)
716 DO k = bo(1, 3), bo(2, 3)
717 DO j = bo(1, 2), bo(2, 2)
718 DO i = bo(1, 1), bo(2, 1)
719 local_row = i - bo(1, 1) + 1 + &
720 nx*(j - bo(1, 2) + ny*(k - bo(1, 3)))
721 grid_point = grid_coordinate(pw_grid, [i, j, k])
722 atomic_base_weight = pw_grid%dvol
723 base_weight = pw_grid%dvol
724 IF (PRESENT(weights)) THEN
725 IF (ASSOCIATED(weights)) base_weight = base_weight*weights%array(i, j, k)
726 END IF
727 cached_layout%feature_index(i, j, k) = local_row
728 IF (nonperiodic) THEN
729 CALL smooth_atom_partition(grid_point, atom_coords_pbc, cell, &
730 partition_weights, partition_atom_coords, distances, &
731 pair_distances)
732 ELSE
733 CALL smooth_atom_partition(grid_point, atom_coords_pbc, cell, &
734 partition_weights, partition_atom_coords, distances)
735 END IF
736 included_sum = sum(partition_weights, mask=partition_weights > smooth_partition_eps)
737 IF (included_sum <= 0.0_dp) THEN
738 owner = nearest_atom(grid_point, atom_coords_pbc, cell)
739 partition_weights = 0.0_dp
740 partition_weights(owner) = 1.0_dp
741 included_sum = 1.0_dp
742 END IF
743 local_feature = 0
744 DO iatom = 1, natom
745 IF (partition_weights(iatom) <= smooth_partition_eps) cycle
746 partition_weight = partition_weights(iatom)/included_sum
747 local_feature = local_feature + 1
748 feature_slot = natom*(local_row - 1) + local_feature
749 local_owner(feature_slot) = iatom
750 local_source_points(feature_slot) = local_row
751 static_base = nstatic_per_point*(feature_slot - 1)
752 local_static(static_base + 1:static_base + 3) = grid_point
753 local_static(static_base + 4) = base_weight*partition_weight
754 local_static(static_base + 5) = atomic_base_weight* &
755 smooth_partition_atomic_weight_scale( &
756 partition_weight)
757 END DO
758 local_feature_counts_tmp(local_row) = local_feature
759 END DO
760 END DO
761 END DO
762!$OMP END DO
763 DEALLOCATE (distances, partition_atom_coords, partition_weights)
764!$OMP END PARALLEL
765
766 ! Compact the conflict-free row slots while preserving the serial row/atom order.
767 nfeature_local = 0
768 DO local_row = 1, nflat_local
769 cpassert(local_feature_counts_tmp(local_row) > 0)
770 DO local_feature = 1, local_feature_counts_tmp(local_row)
771 nfeature_local = nfeature_local + 1
772 feature_slot = natom*(local_row - 1) + local_feature
773 IF (feature_slot == nfeature_local) cycle
774 local_owner(nfeature_local) = local_owner(feature_slot)
775 local_source_points(nfeature_local) = local_source_points(feature_slot)
776 local_static(nstatic_per_point*(nfeature_local - 1) + 1: &
777 nstatic_per_point*nfeature_local) = &
778 local_static(nstatic_per_point*(feature_slot - 1) + 1: &
779 nstatic_per_point*feature_slot)
780 END DO
781 END DO
782 END IF
783 CALL timestop(phase_handle)
784
785 ! SKALA groups all grid points by atom. This ordering is static while the
786 ! grid, cell, atom positions, and optional integration weights are unchanged.
787 CALL timeset("skala_gpw_layout_gather", phase_handle)
788 CALL pw_grid%para%group%allgather(nflat_local, point_counts)
789 CALL counts_to_displs(point_counts, point_displs)
790 npoint = sum(point_counts)
791 CALL pw_grid%para%group%allgather(nfeature_local, feature_counts)
792 CALL counts_to_displs(feature_counts, feature_displs)
793 DO pe = 1, nproc
794 static_counts(pe) = nstatic_per_point*feature_counts(pe)
795 static_displs(pe) = nstatic_per_point*feature_displs(pe)
796 END DO
797 nflat = sum(feature_counts)
798 ALLOCATE (global_owner(nflat), global_source_points(nflat), &
799 global_static(nstatic_per_point*nflat), local_source_global(nfeature_local))
800 DO feature_local = 1, nfeature_local
801 local_source_global(feature_local) = point_displs(pe_index) + local_source_points(feature_local)
802 END DO
803 CALL pw_grid%para%group%allgatherv(local_owner(1:nfeature_local), global_owner, feature_counts, &
804 feature_displs)
805 CALL pw_grid%para%group%allgatherv(local_source_global, global_source_points, feature_counts, &
806 feature_displs)
807 CALL pw_grid%para%group%allgatherv(local_static(1:nstatic_per_point*nfeature_local), &
808 global_static, static_counts, &
809 static_displs)
810 CALL timestop(phase_handle)
811
812 ALLOCATE (cached_layout%chunk_feature_counts(nproc), &
813 cached_layout%chunk_feature_displs(nproc), &
814 cached_layout%chunk_grad_counts(nproc), cached_layout%chunk_grad_displs(nproc), &
815 cached_layout%feature_counts(nproc), cached_layout%feature_displs(nproc), &
816 cached_layout%dynamic_counts(nproc), cached_layout%dynamic_displs(nproc), &
817 cached_layout%route_point_recv_counts(nproc), &
818 cached_layout%route_point_recv_displs(nproc), &
819 cached_layout%route_point_send_counts(nproc), &
820 cached_layout%route_point_send_displs(nproc), &
821 cached_layout%feature_source_points(nflat), &
822 cached_layout%global_to_feature(npoint), cached_layout%atomic_grid_sizes(natom), &
823 cached_layout%local_feature_counts(nflat_local), &
824 cached_layout%local_feature_offsets(nflat_local + 1), &
825 cached_layout%local_feature_rows(nfeature_local), &
826 cached_layout%local_feature_points(nfeature_local), &
827 cached_layout%local_feature_indices(nfeature_local), atom_offset(natom + 1), &
828 atom_position(natom), chunk_atom_begin(nproc), chunk_atom_end(nproc), &
829 cursor(nflat_local))
830 cached_layout%feature_counts(:) = feature_counts
831 cached_layout%feature_displs(:) = feature_displs
832 cached_layout%dynamic_counts(:) = ndynamic_per_point*point_counts
833 cached_layout%dynamic_displs(:) = ndynamic_per_point*point_displs
834 cached_layout%atomic_grid_sizes = 0_int_8
835 cached_layout%global_to_feature = 0
836 cached_layout%local_feature_counts(:) = local_feature_counts_tmp
837 cached_layout%local_feature_offsets(1) = 1
838 DO local_row = 1, nflat_local
839 cached_layout%local_feature_offsets(local_row + 1) = &
840 cached_layout%local_feature_offsets(local_row) + &
841 cached_layout%local_feature_counts(local_row)
842 END DO
843 cursor(:) = cached_layout%local_feature_offsets(1:nflat_local)
844
845 CALL timeset("skala_gpw_layout_atom_sort", phase_handle)
846 DO ipt = 1, nflat
847 cached_layout%atomic_grid_sizes(global_owner(ipt)) = &
848 cached_layout%atomic_grid_sizes(global_owner(ipt)) + 1_int_8
849 END DO
850 atom_offset(1) = 1
851 DO iatom = 1, natom
852 atom_offset(iatom + 1) = atom_offset(iatom) + int(cached_layout%atomic_grid_sizes(iatom))
853 END DO
854 DO iatom = 1, natom
855 atom_position(iatom) = atom_offset(iatom)
856 END DO
857 max_grid_size = maxval(int(cached_layout%atomic_grid_sizes))
858 CALL build_atom_chunks(cached_layout%atomic_grid_sizes, atom_offset, nproc, &
859 chunk_atom_begin, chunk_atom_end, &
860 cached_layout%chunk_feature_counts, &
861 cached_layout%chunk_feature_displs)
862 cached_layout%chunk_grad_counts(:) = ngrad_per_point*cached_layout%chunk_feature_counts
863 cached_layout%chunk_grad_displs(:) = ngrad_per_point*cached_layout%chunk_feature_displs
864 cached_layout%chunk_atom_begin = chunk_atom_begin(pe_index)
865 cached_layout%chunk_atom_end = chunk_atom_end(pe_index)
866 cached_layout%chunk_feature_begin = cached_layout%chunk_feature_displs(pe_index) + 1
867 cached_layout%chunk_feature_count = cached_layout%chunk_feature_counts(pe_index)
868 cached_layout%chunk_natom = cached_layout%chunk_atom_end - &
869 cached_layout%chunk_atom_begin + 1
870
871 ALLOCATE (cached_layout%grid_coords(3, nflat), cached_layout%grid_weights(nflat), &
872 cached_layout%atomic_grid_weights(nflat), &
873 cached_layout%coarse_0_atomic_coords(3, natom), &
874 cached_layout%atomic_grid_size_bound_shape(0, max_grid_size), &
875 cached_layout%atom_coords(3, natom))
876 cached_layout%grid_coords = 0.0_dp
877 cached_layout%grid_weights = 0.0_dp
878 cached_layout%atomic_grid_weights = 0.0_dp
879 cached_layout%atomic_grid_size_bound_shape = 0_int_8
880
881 DO iatom = 1, natom
882 cached_layout%atom_coords(:, iatom) = particle_set(iatom)%r
883 cached_layout%coarse_0_atomic_coords(:, iatom) = atom_coords_pbc(:, iatom)
884 END DO
885
886 DO ipt = 1, nflat
887 owner = global_owner(ipt)
888 row = atom_position(owner)
889 atom_position(owner) = atom_position(owner) + 1
890 source_global = global_source_points(ipt)
891 cached_layout%feature_source_points(row) = source_global
892 IF (cached_layout%global_to_feature(source_global) == 0) THEN
893 cached_layout%global_to_feature(source_global) = row
894 END IF
895 static_base = nstatic_per_point*(ipt - 1)
896 cached_layout%grid_coords(:, row) = global_static(static_base + 1:static_base + 3)
897 cached_layout%grid_weights(row) = global_static(static_base + 4)
898 cached_layout%atomic_grid_weights(row) = global_static(static_base + 5)
899 source_local = source_global - point_displs(pe_index)
900 IF (source_local >= 1 .AND. source_local <= nflat_local) THEN
901 feature_local = cursor(source_local)
902 cursor(source_local) = cursor(source_local) + 1
903 cached_layout%local_feature_rows(feature_local) = row
904 cached_layout%local_feature_points(feature_local) = source_local
905 END IF
906 END DO
907
908 cpassert(all(cached_layout%global_to_feature > 0))
909 cpassert(all(cached_layout%local_feature_rows > 0))
910 cpassert(all(cached_layout%local_feature_points > 0))
911 DO k = bo(1, 3), bo(2, 3)
912 DO j = bo(1, 2), bo(2, 2)
913 DO i = bo(1, 1), bo(2, 1)
914 local_row = cached_layout%feature_index(i, j, k)
915 cached_layout%feature_index(i, j, k) = &
916 cached_layout%local_feature_rows(cached_layout%local_feature_offsets(local_row))
917 END DO
918 END DO
919 END DO
920 DO feature_local = 1, nfeature_local
921 cached_layout%local_feature_indices(feature_local) = &
922 int(cached_layout%local_feature_rows(feature_local) - 1, kind=int_8)
923 END DO
924 CALL timestop(phase_handle)
925 CALL timeset("skala_gpw_layout_chunk_routes", phase_handle)
926 CALL build_atom_chunk_routes(cached_layout, cached_layout%local_feature_rows, &
927 pw_grid%para%group)
928 CALL build_atom_chunk_layout(cached_layout)
929 CALL timestop(phase_handle)
930
931 cached_layout%natom = natom
932 cached_layout%nflat = nflat
933 cached_layout%nflat_local = nflat_local
934 cached_layout%npoint = npoint
935 cached_layout%nproc = nproc
936 cached_layout%atom_partition = my_atom_partition
937 cached_layout%bo = bo
938 cached_layout%bounds = pw_grid%bounds
939 cached_layout%npts = pw_grid%npts
940 cached_layout%dvol = pw_grid%dvol
941 cached_layout%dh = pw_grid%dh
942 cached_layout%cell_hmat = cell%hmat
943 cached_layout%weight_sum = weight_sum
944 cached_layout%weight_sumsq = weight_sumsq
945 cached_layout%has_weights = has_weights
946 CALL timeset("skala_gpw_layout_tensors", phase_handle)
947 IF (needs_full_static_tensors) CALL build_full_static_layout_tensors(cached_layout)
948 IF (cached_layout%chunk_feature_count > 0) THEN
949 CALL build_chunk_static_layout_tensors(cached_layout)
950 END IF
951 CALL timestop(phase_handle)
952 cached_layout%active = .true.
953
954 DEALLOCATE (atom_coords_pbc, atom_offset, atom_position, &
955 chunk_atom_begin, chunk_atom_end, cursor, feature_counts, feature_displs, &
956 global_owner, global_source_points, global_static, local_feature_counts_tmp, &
957 local_owner, local_source_global, local_source_points, &
958 local_static, pair_distances, point_counts, point_displs, static_counts, &
959 static_displs)
960
961 END SUBROUTINE rebuild_layout_cache
962
963! **************************************************************************************************
964!> \brief Build cached full-layout Torch tensors for static SKALA inputs.
965!> \param cache ...
966! **************************************************************************************************
967 SUBROUTINE build_full_static_layout_tensors(cache)
968 TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
969
970 cpassert(.NOT. cache%static_tensors_active)
971
972 CALL torch_tensor_from_array(cache%grid_coords_t, cache%grid_coords)
973 CALL torch_tensor_to_device_leaf(cache%grid_coords_t, .false.)
974 CALL torch_tensor_from_array(cache%grid_weights_t, cache%grid_weights)
975 CALL torch_tensor_to_device_leaf(cache%grid_weights_t, .false.)
976 CALL torch_tensor_from_array(cache%atomic_grid_weights_t, cache%atomic_grid_weights)
977 CALL torch_tensor_to_device_leaf(cache%atomic_grid_weights_t, .false.)
978 CALL torch_tensor_from_array(cache%atomic_grid_sizes_t, cache%atomic_grid_sizes)
979 CALL torch_tensor_to_device_leaf(cache%atomic_grid_sizes_t, .false.)
980 CALL torch_tensor_from_array(cache%coarse_0_atomic_coords_t, cache%coarse_0_atomic_coords)
981 CALL torch_tensor_to_device_leaf(cache%coarse_0_atomic_coords_t, .false.)
982 CALL torch_tensor_from_array(cache%atomic_grid_size_bound_shape_t, &
983 cache%atomic_grid_size_bound_shape)
984 CALL torch_tensor_to_device_leaf(cache%atomic_grid_size_bound_shape_t, .false.)
985 CALL torch_tensor_from_array(cache%local_feature_indices_t, cache%local_feature_indices)
986 CALL torch_tensor_to_device_leaf(cache%local_feature_indices_t, .false.)
987
988 CALL torch_dict_create(cache%static_inputs)
989 CALL torch_dict_insert(cache%static_inputs, "grid_coords", cache%grid_coords_t)
990 CALL torch_dict_insert(cache%static_inputs, "grid_weights", cache%grid_weights_t)
991 CALL torch_dict_insert(cache%static_inputs, "atomic_grid_weights", &
992 cache%atomic_grid_weights_t)
993 CALL torch_dict_insert(cache%static_inputs, "atomic_grid_sizes", &
994 cache%atomic_grid_sizes_t)
995 CALL torch_dict_insert(cache%static_inputs, "atomic_grid_size_bound_shape", &
996 cache%atomic_grid_size_bound_shape_t)
997 cache%static_tensors_active = .true.
998
999 END SUBROUTINE build_full_static_layout_tensors
1000
1001! **************************************************************************************************
1002!> \brief Build cached atom-chunk Torch tensors for static SKALA inputs.
1003!> \param cache ...
1004! **************************************************************************************************
1005 SUBROUTINE build_chunk_static_layout_tensors(cache)
1006 TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1007
1008 cpassert(.NOT. cache%chunk_static_tensors_active)
1009 CALL torch_tensor_from_array(cache%chunk_grid_coords_t, cache%chunk_grid_coords)
1010 CALL torch_tensor_to_device_leaf(cache%chunk_grid_coords_t, .false.)
1011 CALL torch_tensor_from_array(cache%chunk_grid_weights_t, cache%chunk_grid_weights)
1012 CALL torch_tensor_to_device_leaf(cache%chunk_grid_weights_t, .false.)
1013 CALL torch_tensor_from_array(cache%chunk_atomic_grid_weights_t, &
1014 cache%chunk_atomic_grid_weights)
1015 CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_weights_t, .false.)
1016 CALL torch_tensor_from_array(cache%chunk_atomic_grid_sizes_t, &
1017 cache%chunk_atomic_grid_sizes)
1018 CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_sizes_t, .false.)
1019 CALL torch_tensor_from_array(cache%chunk_coarse_0_atomic_coords_t, &
1020 cache%chunk_coarse_0_atomic_coords)
1021 CALL torch_tensor_to_device_leaf(cache%chunk_coarse_0_atomic_coords_t, .false.)
1022 CALL torch_tensor_from_array(cache%chunk_atomic_grid_size_bound_shape_t, &
1023 cache%chunk_atomic_grid_size_bound_shape)
1024 CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_size_bound_shape_t, .false.)
1025 CALL torch_tensor_from_array(cache%chunk_feature_indices_t, cache%chunk_feature_indices)
1026 CALL torch_tensor_to_device_leaf(cache%chunk_feature_indices_t, .false.)
1027
1028 CALL torch_dict_create(cache%chunk_static_inputs)
1029 CALL torch_dict_insert(cache%chunk_static_inputs, "grid_coords", &
1030 cache%chunk_grid_coords_t)
1031 CALL torch_dict_insert(cache%chunk_static_inputs, "grid_weights", &
1032 cache%chunk_grid_weights_t)
1033 CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_weights", &
1034 cache%chunk_atomic_grid_weights_t)
1035 CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_sizes", &
1036 cache%chunk_atomic_grid_sizes_t)
1037 CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_size_bound_shape", &
1038 cache%chunk_atomic_grid_size_bound_shape_t)
1039 cache%chunk_static_tensors_active = .true.
1040
1041 END SUBROUTINE build_chunk_static_layout_tensors
1042
1043! **************************************************************************************************
1044!> \brief Copy static cached layout arrays into a feature bundle.
1045!> \param features ...
1046!> \param use_atom_chunk_routing ...
1047!> \param needs_coordinate_array ...
1048!> \param needs_grid_coordinate_array ...
1049! **************************************************************************************************
1050 SUBROUTINE copy_cached_layout(features, use_atom_chunk_routing, needs_coordinate_array, &
1051 needs_grid_coordinate_array)
1052 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1053 LOGICAL, INTENT(IN) :: use_atom_chunk_routing, &
1054 needs_coordinate_array, &
1055 needs_grid_coordinate_array
1056
1057 cpassert(cached_layout%active)
1058
1059 IF (.NOT. use_atom_chunk_routing) THEN
1060 ALLOCATE (features%feature_index(lbound(cached_layout%feature_index, 1): &
1061 ubound(cached_layout%feature_index, 1), &
1062 lbound(cached_layout%feature_index, 2): &
1063 ubound(cached_layout%feature_index, 2), &
1064 lbound(cached_layout%feature_index, 3): &
1065 ubound(cached_layout%feature_index, 3)))
1066 ALLOCATE (features%grid_weights(cached_layout%nflat))
1067 ALLOCATE (features%local_feature_counts(cached_layout%nflat_local), &
1068 features%local_feature_offsets(cached_layout%nflat_local + 1), &
1069 features%local_feature_rows(SIZE(cached_layout%local_feature_rows)))
1070
1071 features%feature_index(:, :, :) = cached_layout%feature_index
1072 features%grid_weights(:) = cached_layout%grid_weights
1073 features%local_feature_counts(:) = cached_layout%local_feature_counts
1074 features%local_feature_offsets(:) = cached_layout%local_feature_offsets
1075 features%local_feature_rows(:) = cached_layout%local_feature_rows
1076 END IF
1077 features%nflat = cached_layout%nflat
1078 features%nflat_local = cached_layout%nflat_local
1079 features%chunk_feature_count = cached_layout%chunk_feature_count
1080 features%atom_partition = cached_layout%atom_partition
1081 ALLOCATE (features%atomic_grid_sizes(cached_layout%natom))
1082 features%atomic_grid_sizes(:) = cached_layout%atomic_grid_sizes
1083 IF (needs_grid_coordinate_array) THEN
1084 ALLOCATE (features%grid_coords(3, cached_layout%nflat))
1085 ALLOCATE (features%atomic_grid_weights(cached_layout%nflat))
1086 features%grid_coords(:, :) = cached_layout%grid_coords
1087 features%atomic_grid_weights(:) = cached_layout%atomic_grid_weights
1088 END IF
1089 ALLOCATE (features%chunk_grad_counts(cached_layout%nproc), &
1090 features%chunk_grad_displs(cached_layout%nproc))
1091 features%chunk_grad_counts(:) = cached_layout%chunk_grad_counts
1092 features%chunk_grad_displs(:) = cached_layout%chunk_grad_displs
1093 IF (use_atom_chunk_routing) THEN
1094 features%route_point_recv_counts => cached_layout%route_point_recv_counts
1095 features%route_point_recv_displs => cached_layout%route_point_recv_displs
1096 features%route_point_send_counts => cached_layout%route_point_send_counts
1097 features%route_point_send_displs => cached_layout%route_point_send_displs
1098 features%route_chunk_offsets => cached_layout%route_chunk_offsets
1099 features%route_chunk_rows => cached_layout%route_chunk_rows
1100 features%route_row_offsets => cached_layout%route_row_offsets
1101 features%route_row_positions => cached_layout%route_row_positions
1102 features%route_send_local_rows => cached_layout%route_send_local_rows
1103 END IF
1104 IF (needs_coordinate_array) THEN
1105 ALLOCATE (features%coarse_0_atomic_coords(3, cached_layout%natom))
1106 features%coarse_0_atomic_coords(:, :) = cached_layout%coarse_0_atomic_coords
1107 END IF
1108
1109 END SUBROUTINE copy_cached_layout
1110
1111! **************************************************************************************************
1112!> \brief Split the atom-ordered feature rows into contiguous atom chunks.
1113!> \param atomic_grid_sizes ...
1114!> \param atom_offset ...
1115!> \param nproc ...
1116!> \param chunk_atom_begin ...
1117!> \param chunk_atom_end ...
1118!> \param chunk_feature_counts ...
1119!> \param chunk_feature_displs ...
1120! **************************************************************************************************
1121 SUBROUTINE build_atom_chunks(atomic_grid_sizes, atom_offset, nproc, chunk_atom_begin, &
1122 chunk_atom_end, chunk_feature_counts, chunk_feature_displs)
1123 INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1124 INTEGER, DIMENSION(:), INTENT(IN) :: atom_offset
1125 INTEGER, INTENT(IN) :: nproc
1126 INTEGER, DIMENSION(:), INTENT(OUT) :: chunk_atom_begin, chunk_atom_end, &
1127 chunk_feature_counts, &
1128 chunk_feature_displs
1129
1130 INTEGER :: best_limit, count, displ, end_atom, lower_limit, max_end_atom, midpoint, natom, &
1131 next_atom, next_count, pe, ranks_left, target_chunks, total_count, upper_limit
1132
1133 natom = SIZE(atomic_grid_sizes)
1134 chunk_atom_begin = natom + 1
1135 chunk_atom_end = natom
1136 chunk_feature_counts = 0
1137 chunk_feature_displs = 0
1138 IF (natom == 0) RETURN
1139
1140 target_chunks = min(nproc, natom)
1141 total_count = atom_offset(natom + 1) - 1
1142 lower_limit = maxval(int(atomic_grid_sizes))
1143 lower_limit = max(lower_limit, (total_count + target_chunks - 1)/target_chunks)
1144 upper_limit = total_count
1145 best_limit = upper_limit
1146 DO WHILE (lower_limit <= upper_limit)
1147 midpoint = (lower_limit + upper_limit)/2
1148 IF (atom_chunks_fit_limit(atomic_grid_sizes, midpoint, target_chunks)) THEN
1149 best_limit = midpoint
1150 upper_limit = midpoint - 1
1151 ELSE
1152 lower_limit = midpoint + 1
1153 END IF
1154 END DO
1155
1156 displ = 0
1157 next_atom = 1
1158 DO pe = 1, nproc
1159 chunk_feature_displs(pe) = displ
1160 IF (pe > target_chunks .OR. next_atom > natom) cycle
1161
1162 ranks_left = target_chunks - pe + 1
1163 chunk_atom_begin(pe) = next_atom
1164 max_end_atom = natom - ranks_left + 1
1165 end_atom = next_atom
1166 count = int(atomic_grid_sizes(end_atom))
1167 DO WHILE (end_atom < max_end_atom)
1168 next_count = count + int(atomic_grid_sizes(end_atom + 1))
1169 IF (next_count > best_limit) EXIT
1170 end_atom = end_atom + 1
1171 count = next_count
1172 END DO
1173
1174 chunk_atom_end(pe) = end_atom
1175 chunk_feature_counts(pe) = atom_offset(end_atom + 1) - atom_offset(next_atom)
1176 displ = displ + chunk_feature_counts(pe)
1177 next_atom = end_atom + 1
1178 END DO
1179
1180 cpassert(displ == atom_offset(natom + 1) - 1)
1181
1182 END SUBROUTINE build_atom_chunks
1183
1184! **************************************************************************************************
1185!> \brief Check if contiguous atom chunks can stay below a feature-count limit.
1186!> \param atomic_grid_sizes ...
1187!> \param limit ...
1188!> \param nchunks ...
1189!> \return ...
1190! **************************************************************************************************
1191 FUNCTION atom_chunks_fit_limit(atomic_grid_sizes, limit, nchunks) RESULT(fits)
1192 INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1193 INTEGER, INTENT(IN) :: limit, nchunks
1194 LOGICAL :: fits
1195
1196 INTEGER :: atom_count, chunk_count, iatom, &
1197 used_chunks
1198
1199 fits = .false.
1200 IF (SIZE(atomic_grid_sizes) == 0) THEN
1201 fits = .true.
1202 RETURN
1203 END IF
1204
1205 used_chunks = 1
1206 chunk_count = 0
1207 DO iatom = 1, SIZE(atomic_grid_sizes)
1208 atom_count = int(atomic_grid_sizes(iatom))
1209 IF (atom_count > limit) RETURN
1210 IF (chunk_count + atom_count > limit) THEN
1211 used_chunks = used_chunks + 1
1212 chunk_count = atom_count
1213 ELSE
1214 chunk_count = chunk_count + atom_count
1215 END IF
1216 END DO
1217 fits = used_chunks <= nchunks
1218
1219 END FUNCTION atom_chunks_fit_limit
1220
1221! **************************************************************************************************
1222!> \brief Return the MPI rank owning an atom-ordered feature row.
1223!> \param row ...
1224!> \param counts ...
1225!> \param displs ...
1226!> \return ...
1227! **************************************************************************************************
1228 FUNCTION feature_row_chunk_owner(row, counts, displs) RESULT(owner)
1229 INTEGER, INTENT(IN) :: row
1230 INTEGER, DIMENSION(:), INTENT(IN) :: counts, displs
1231 INTEGER :: owner
1232
1233 INTEGER :: pe
1234
1235 owner = 0
1236 DO pe = 1, SIZE(counts)
1237 IF (row > displs(pe) .AND. row <= displs(pe) + counts(pe)) THEN
1238 owner = pe
1239 RETURN
1240 END IF
1241 END DO
1242
1243 END FUNCTION feature_row_chunk_owner
1244
1245! **************************************************************************************************
1246!> \brief Build zero-based displacement arrays from per-rank counts.
1247!> \param counts ...
1248!> \param displs ...
1249! **************************************************************************************************
1250 SUBROUTINE counts_to_displs(counts, displs)
1251 INTEGER, DIMENSION(:), INTENT(IN) :: counts
1252 INTEGER, DIMENSION(:), INTENT(OUT) :: displs
1253
1254 INTEGER :: pe
1255
1256 displs(1) = 0
1257 DO pe = 2, SIZE(counts)
1258 displs(pe) = displs(pe - 1) + counts(pe - 1)
1259 END DO
1260
1261 END SUBROUTINE counts_to_displs
1262
1263! **************************************************************************************************
1264!> \brief Precompute all-to-all routing between local grid rows and atom chunks.
1265!> \param cache ...
1266!> \param local_to_global ...
1267!> \param group ...
1268! **************************************************************************************************
1269 SUBROUTINE build_atom_chunk_routes(cache, local_to_global, group)
1270 TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1271 INTEGER, DIMENSION(:), INTENT(IN) :: local_to_global
1272
1273 CLASS(mp_comm_type), INTENT(IN) :: group
1274
1275 INTEGER :: chunk_row, dest, feature_begin, feature_end, &
1276 feature_pos, local_feature, local_row, meta_pos, &
1277 nflat_local, nproc, nrecv, nsend, pe, point_pos, &
1278 route_index, row, row_position
1279 INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_cursor, chunk_route_counts, &
1280 feature_cursor, feature_recv_counts, &
1281 feature_recv_displs, feature_send_counts, &
1282 feature_send_displs, meta_recv_counts, &
1283 meta_recv_displs, meta_send_counts, &
1284 meta_send_displs, recv_meta, route_cursor, &
1285 route_feature_dest, route_last_row, &
1286 route_last_position, send_meta
1287
1288 nflat_local = SIZE(cache%local_feature_offsets) - 1
1289 nproc = SIZE(cache%route_point_send_counts)
1290 ALLOCATE (route_feature_dest(SIZE(local_to_global)), &
1291 feature_send_counts(nproc), feature_send_displs(nproc), &
1292 feature_recv_counts(nproc), feature_recv_displs(nproc), &
1293 route_last_row(nproc), route_last_position(nproc))
1294 cache%route_point_send_counts = 0
1295 feature_send_counts = 0
1296 route_feature_dest = 0
1297 route_last_row = 0
1298 ! Route each local grid row only once to each atom-chunk owner.
1299 DO local_row = 1, nflat_local
1300 feature_begin = cache%local_feature_offsets(local_row)
1301 feature_end = cache%local_feature_offsets(local_row + 1) - 1
1302 DO local_feature = feature_begin, feature_end
1303 dest = feature_row_chunk_owner(local_to_global(local_feature), &
1304 cache%chunk_feature_counts, &
1305 cache%chunk_feature_displs)
1306 cpassert(dest > 0)
1307 route_feature_dest(local_feature) = dest
1308 feature_send_counts(dest) = feature_send_counts(dest) + 1
1309 IF (route_last_row(dest) /= local_row) THEN
1310 cache%route_point_send_counts(dest) = cache%route_point_send_counts(dest) + 1
1311 route_last_row(dest) = local_row
1312 END IF
1313 END DO
1314 END DO
1315 CALL counts_to_displs(cache%route_point_send_counts, cache%route_point_send_displs)
1316 CALL counts_to_displs(feature_send_counts, feature_send_displs)
1317 CALL group%alltoall(cache%route_point_send_counts, cache%route_point_recv_counts, 1)
1318 CALL counts_to_displs(cache%route_point_recv_counts, cache%route_point_recv_displs)
1319 CALL group%alltoall(feature_send_counts, feature_recv_counts, 1)
1320 CALL counts_to_displs(feature_recv_counts, feature_recv_displs)
1321
1322 nsend = sum(cache%route_point_send_counts)
1323 nrecv = sum(cache%route_point_recv_counts)
1324 ALLOCATE (cache%route_send_local_rows(nsend), &
1325 cache%route_row_offsets(nflat_local + 1), &
1326 cache%route_row_positions(nsend), &
1327 route_cursor(nproc), feature_cursor(nproc), &
1328 send_meta(2*SIZE(local_to_global)), &
1329 recv_meta(2*cache%chunk_feature_count), &
1330 meta_send_counts(nproc), meta_send_displs(nproc), &
1331 meta_recv_counts(nproc), meta_recv_displs(nproc))
1332 cache%route_send_local_rows = 0
1333 cache%route_row_positions = 0
1334 send_meta = 0
1335 route_cursor(:) = cache%route_point_send_displs + 1
1336 feature_cursor(:) = feature_send_displs + 1
1337 route_last_row = 0
1338 route_last_position = 0
1339 cache%route_row_offsets(1) = 1
1340 row_position = 1
1341 DO local_row = 1, nflat_local
1342 feature_begin = cache%local_feature_offsets(local_row)
1343 feature_end = cache%local_feature_offsets(local_row + 1) - 1
1344 DO local_feature = feature_begin, feature_end
1345 dest = route_feature_dest(local_feature)
1346 IF (route_last_row(dest) /= local_row) THEN
1347 point_pos = route_cursor(dest)
1348 route_cursor(dest) = route_cursor(dest) + 1
1349 route_last_row(dest) = local_row
1350 route_last_position(dest) = point_pos
1351 cache%route_send_local_rows(point_pos) = local_row
1352 cache%route_row_positions(row_position) = point_pos
1353 row_position = row_position + 1
1354 ELSE
1355 point_pos = route_last_position(dest)
1356 END IF
1357 feature_pos = feature_cursor(dest)
1358 feature_cursor(dest) = feature_cursor(dest) + 1
1359 send_meta(2*feature_pos - 1) = local_to_global(local_feature)
1360 send_meta(2*feature_pos) = point_pos - cache%route_point_send_displs(dest)
1361 END DO
1362 cache%route_row_offsets(local_row + 1) = row_position
1363 END DO
1364 cpassert(row_position == nsend + 1)
1365
1366 ! Tell each chunk owner which feature rows share a compressed route point.
1367 meta_send_counts(:) = 2*feature_send_counts
1368 meta_send_displs(:) = 2*feature_send_displs
1369 meta_recv_counts(:) = 2*feature_recv_counts
1370 meta_recv_displs(:) = 2*feature_recv_displs
1371 CALL group%alltoall(send_meta, meta_send_counts, meta_send_displs, recv_meta, &
1372 meta_recv_counts, meta_recv_displs)
1373
1374 ! Store the route-to-feature expansion in CSR form for gradient reduction.
1375 ALLOCATE (cache%chunk_return_positions(cache%chunk_feature_count), &
1376 cache%route_chunk_offsets(nrecv + 1), &
1377 cache%route_chunk_rows(cache%chunk_feature_count), &
1378 chunk_route_counts(nrecv), chunk_cursor(nrecv))
1379 cache%chunk_return_positions = 0
1380 cache%route_chunk_rows = 0
1381 chunk_route_counts = 0
1382 DO pe = 1, nproc
1383 DO feature_pos = feature_recv_displs(pe) + 1, &
1384 feature_recv_displs(pe) + feature_recv_counts(pe)
1385 meta_pos = 2*feature_pos
1386 row = recv_meta(meta_pos - 1)
1387 route_index = recv_meta(meta_pos)
1388 point_pos = cache%route_point_recv_displs(pe) + route_index
1389 cpassert(point_pos >= 1 .AND. point_pos <= nrecv)
1390 chunk_row = row - cache%chunk_feature_begin + 1
1391 cpassert(chunk_row >= 1 .AND. chunk_row <= cache%chunk_feature_count)
1392 cache%chunk_return_positions(chunk_row) = point_pos
1393 chunk_route_counts(point_pos) = chunk_route_counts(point_pos) + 1
1394 END DO
1395 END DO
1396 cache%route_chunk_offsets(1) = 1
1397 DO point_pos = 1, nrecv
1398 cache%route_chunk_offsets(point_pos + 1) = &
1399 cache%route_chunk_offsets(point_pos) + chunk_route_counts(point_pos)
1400 END DO
1401 chunk_cursor(:) = cache%route_chunk_offsets(1:nrecv)
1402 DO pe = 1, nproc
1403 DO feature_pos = feature_recv_displs(pe) + 1, &
1404 feature_recv_displs(pe) + feature_recv_counts(pe)
1405 meta_pos = 2*feature_pos
1406 row = recv_meta(meta_pos - 1)
1407 chunk_row = row - cache%chunk_feature_begin + 1
1408 point_pos = cache%chunk_return_positions(chunk_row)
1409 cache%route_chunk_rows(chunk_cursor(point_pos)) = chunk_row
1410 chunk_cursor(point_pos) = chunk_cursor(point_pos) + 1
1411 END DO
1412 END DO
1413
1414 cpassert(sum(feature_send_counts) == SIZE(local_to_global))
1415 cpassert(sum(feature_recv_counts) == cache%chunk_feature_count)
1416 cpassert(sum(cache%route_point_send_counts) == nsend)
1417 cpassert(sum(cache%route_point_recv_counts) == nrecv)
1418 cpassert(all(cache%chunk_return_positions > 0))
1419 cpassert(all(cache%route_chunk_rows > 0))
1420 cpassert(all(cache%route_row_positions > 0))
1421 cpassert(all(cache%route_send_local_rows > 0))
1422
1423 DEALLOCATE (chunk_cursor, chunk_route_counts, feature_cursor, feature_recv_counts, &
1424 feature_recv_displs, feature_send_counts, feature_send_displs, &
1425 meta_recv_counts, meta_recv_displs, meta_send_counts, meta_send_displs, &
1426 recv_meta, route_cursor, route_feature_dest, route_last_position, route_last_row, &
1427 send_meta)
1428
1429 END SUBROUTINE build_atom_chunk_routes
1430
1431! **************************************************************************************************
1432!> \brief Materialize the current rank's atom chunk static layout.
1433!> \param cache ...
1434! **************************************************************************************************
1435 SUBROUTINE build_atom_chunk_layout(cache)
1436 TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1437
1438 INTEGER :: irow, max_grid_size, row_begin, row_end
1439
1440 IF (cache%chunk_feature_count <= 0 .OR. cache%chunk_natom <= 0) RETURN
1441
1442 row_begin = cache%chunk_feature_begin
1443 row_end = row_begin + cache%chunk_feature_count - 1
1444 ALLOCATE (cache%chunk_grid_coords(3, cache%chunk_feature_count), &
1445 cache%chunk_grid_weights(cache%chunk_feature_count), &
1446 cache%chunk_atomic_grid_weights(cache%chunk_feature_count), &
1447 cache%chunk_atomic_grid_sizes(cache%chunk_natom), &
1448 cache%chunk_coarse_0_atomic_coords(3, cache%chunk_natom), &
1449 cache%chunk_feature_indices(cache%chunk_feature_count))
1450 cache%chunk_grid_coords(:, :) = cache%grid_coords(:, row_begin:row_end)
1451 cache%chunk_grid_weights(:) = cache%grid_weights(row_begin:row_end)
1452 cache%chunk_atomic_grid_weights(:) = cache%atomic_grid_weights(row_begin:row_end)
1453 cache%chunk_atomic_grid_sizes(:) = &
1454 cache%atomic_grid_sizes(cache%chunk_atom_begin:cache%chunk_atom_end)
1455 cache%chunk_coarse_0_atomic_coords(:, :) = &
1456 cache%coarse_0_atomic_coords(:, cache%chunk_atom_begin:cache%chunk_atom_end)
1457
1458 max_grid_size = maxval(int(cache%chunk_atomic_grid_sizes))
1459 ALLOCATE (cache%chunk_atomic_grid_size_bound_shape(0, max_grid_size))
1460 cache%chunk_atomic_grid_size_bound_shape = 0_int_8
1461 DO irow = 1, cache%chunk_feature_count
1462 cache%chunk_feature_indices(irow) = int(irow - 1, kind=int_8)
1463 END DO
1464
1465 END SUBROUTINE build_atom_chunk_layout
1466
1467! **************************************************************************************************
1468!> \brief Send local dynamic feature rows to their atom-chunk owner ranks.
1469!> \param features ...
1470!> \param local_dynamic ...
1471!> \param group ...
1472!> \param collapse_spin_dynamics ...
1473! **************************************************************************************************
1474 SUBROUTINE route_atom_chunk_dynamics(features, local_dynamic, group, collapse_spin_dynamics)
1475 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1476 REAL(kind=dp), DIMENSION(:), INTENT(IN) :: local_dynamic
1477
1478 CLASS(mp_comm_type), INTENT(IN) :: group
1479 LOGICAL, INTENT(IN) :: collapse_spin_dynamics
1480
1481 INTEGER :: chunk_row, dyn_base, local_row, &
1482 ndynamic_route_per_point, nrecv, nsend, &
1483 point_pos, src_base
1484 INTEGER, ALLOCATABLE, DIMENSION(:) :: recv_counts, recv_displs, send_counts, send_displs
1485 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: recv_dynamic, send_dynamic
1486
1487 nsend = SIZE(cached_layout%route_send_local_rows)
1488 nrecv = sum(cached_layout%route_point_recv_counts)
1489 cpassert(nsend == sum(cached_layout%route_point_send_counts))
1490 ndynamic_route_per_point = ndynamic_per_point
1491 IF (collapse_spin_dynamics) ndynamic_route_per_point = nrks_dynamic_per_point
1492
1493 ALLOCATE (send_dynamic(max(1, ndynamic_route_per_point*nsend)), &
1494 recv_dynamic(max(1, ndynamic_route_per_point*nrecv)), &
1495 send_counts(cached_layout%nproc), send_displs(cached_layout%nproc), &
1496 recv_counts(cached_layout%nproc), &
1497 recv_displs(cached_layout%nproc))
1498 send_counts(:) = ndynamic_route_per_point*cached_layout%route_point_send_counts
1499 send_displs(:) = ndynamic_route_per_point*cached_layout%route_point_send_displs
1500 recv_counts(:) = ndynamic_route_per_point*cached_layout%route_point_recv_counts
1501 recv_displs(:) = ndynamic_route_per_point*cached_layout%route_point_recv_displs
1502!$OMP PARALLEL DO DEFAULT(NONE) &
1503!$OMP SHARED(cached_layout, local_dynamic, ndynamic_route_per_point, nsend, send_dynamic) &
1504!$OMP PRIVATE(dyn_base, local_row, point_pos, src_base)
1505 DO point_pos = 1, nsend
1506 dyn_base = ndynamic_route_per_point*(point_pos - 1)
1507 local_row = cached_layout%route_send_local_rows(point_pos)
1508 src_base = ndynamic_route_per_point*(local_row - 1)
1509 send_dynamic(dyn_base + 1:dyn_base + ndynamic_route_per_point) = &
1510 local_dynamic(src_base + 1:src_base + ndynamic_route_per_point)
1511 END DO
1512!$OMP END PARALLEL DO
1513
1514 CALL group%alltoall(send_dynamic, send_counts, send_displs, recv_dynamic, recv_counts, &
1515 recv_displs)
1516
1517 features%uses_collapsed_rks_dynamic = collapse_spin_dynamics
1518 IF (cached_layout%chunk_feature_count > 0) THEN
1519 IF (collapse_spin_dynamics) THEN
1520 ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 1), &
1521 features%chunk_grad(cached_layout%chunk_feature_count, 3, 1), &
1522 features%chunk_kin(cached_layout%chunk_feature_count, 1))
1523 ELSE
1524 ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1525 features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1526 features%chunk_kin(cached_layout%chunk_feature_count, 2))
1527 END IF
1528
1529!$OMP PARALLEL DO DEFAULT(NONE) &
1530!$OMP SHARED(cached_layout, collapse_spin_dynamics, features, ndynamic_route_per_point, nrecv, recv_dynamic) &
1531!$OMP PRIVATE(chunk_row, dyn_base, point_pos)
1532 DO chunk_row = 1, cached_layout%chunk_feature_count
1533 point_pos = cached_layout%chunk_return_positions(chunk_row)
1534 cpassert(point_pos >= 1 .AND. point_pos <= nrecv)
1535 dyn_base = ndynamic_route_per_point*(point_pos - 1)
1536 IF (collapse_spin_dynamics) THEN
1537 features%chunk_density(chunk_row, 1) = recv_dynamic(dyn_base + 1)
1538 features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 2)
1539 features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 3)
1540 features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 4)
1541 features%chunk_kin(chunk_row, 1) = recv_dynamic(dyn_base + 5)
1542 ELSE
1543 features%chunk_density(chunk_row, :) = recv_dynamic(dyn_base + 1:dyn_base + 2)
1544 features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 3)
1545 features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 4)
1546 features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 5)
1547 features%chunk_grad(chunk_row, 1, 2) = recv_dynamic(dyn_base + 6)
1548 features%chunk_grad(chunk_row, 2, 2) = recv_dynamic(dyn_base + 7)
1549 features%chunk_grad(chunk_row, 3, 2) = recv_dynamic(dyn_base + 8)
1550 features%chunk_kin(chunk_row, :) = recv_dynamic(dyn_base + 9:dyn_base + 10)
1551 END IF
1552 END DO
1553!$OMP END PARALLEL DO
1554 END IF
1555
1556 DEALLOCATE (recv_counts, recv_displs, recv_dynamic, send_counts, send_displs, send_dynamic)
1557
1558 END SUBROUTINE route_atom_chunk_dynamics
1559
1560! **************************************************************************************************
1561!> \brief Extract the current rank's atom chunk from the global dynamic feature arrays.
1562!> \param features ...
1563! **************************************************************************************************
1564 SUBROUTINE extract_atom_chunk_dynamics(features)
1565 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1566
1567 INTEGER :: row_begin, row_end
1568
1569 cpassert(cached_layout%chunk_feature_count > 0)
1570 row_begin = cached_layout%chunk_feature_begin
1571 row_end = row_begin + cached_layout%chunk_feature_count - 1
1572 ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1573 features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1574 features%chunk_kin(cached_layout%chunk_feature_count, 2))
1575 features%chunk_density(:, :) = features%density(row_begin:row_end, :)
1576 features%chunk_grad(:, :, :) = features%grad(row_begin:row_end, :, :)
1577 features%chunk_kin(:, :) = features%kin(row_begin:row_end, :)
1578
1579 END SUBROUTINE extract_atom_chunk_dynamics
1580
1581! **************************************************************************************************
1582!> \brief Compute a local signature for optional integration weights.
1583!> \param weights ...
1584!> \param has_weights ...
1585!> \param weight_sum ...
1586!> \param weight_sumsq ...
1587! **************************************************************************************************
1588 SUBROUTINE weights_signature(weights, has_weights, weight_sum, weight_sumsq)
1589 TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
1590 LOGICAL, INTENT(OUT) :: has_weights
1591 REAL(kind=dp), INTENT(OUT) :: weight_sum, weight_sumsq
1592
1593 has_weights = .false.
1594 weight_sum = 0.0_dp
1595 weight_sumsq = 0.0_dp
1596 IF (PRESENT(weights)) THEN
1597 IF (ASSOCIATED(weights)) THEN
1598 has_weights = .true.
1599 weight_sum = sum(weights%array)
1600 weight_sumsq = sum(weights%array*weights%array)
1601 END IF
1602 END IF
1603
1604 END SUBROUTINE weights_signature
1605
1606! **************************************************************************************************
1607!> \brief Release cached layout arrays.
1608!> \param cache ...
1609! **************************************************************************************************
1610 SUBROUTINE release_layout_cache(cache)
1611 TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1612
1613 IF (cache%inputs_active) THEN
1614 CALL torch_dict_release(cache%inputs)
1615 cache%inputs_active = .false.
1616 END IF
1617
1618 IF (cache%chunk_inputs_active) THEN
1619 CALL torch_dict_release(cache%chunk_inputs)
1620 cache%chunk_inputs_active = .false.
1621 END IF
1622
1623 IF (cache%dynamic_tensors_active) THEN
1624 CALL torch_tensor_release(cache%density_t)
1625 CALL torch_tensor_release(cache%grad_t)
1626 CALL torch_tensor_release(cache%kin_t)
1627 cache%dynamic_tensors_active = .false.
1628 END IF
1629
1630 IF (cache%chunk_dynamic_tensors_active) THEN
1631 IF (cache%chunk_dynamic_input_views_active) THEN
1632 CALL torch_tensor_release(cache%chunk_density_input_t)
1633 CALL torch_tensor_release(cache%chunk_grad_input_t)
1634 CALL torch_tensor_release(cache%chunk_kin_input_t)
1635 cache%chunk_dynamic_input_views_active = .false.
1636 END IF
1637 CALL torch_tensor_release(cache%chunk_density_t)
1638 CALL torch_tensor_release(cache%chunk_grad_t)
1639 CALL torch_tensor_release(cache%chunk_kin_t)
1640 cache%chunk_dynamic_tensors_active = .false.
1641 END IF
1642
1643 IF (cache%static_tensors_active) THEN
1644 CALL torch_tensor_release(cache%grid_coords_t)
1645 CALL torch_tensor_release(cache%grid_weights_t)
1646 CALL torch_tensor_release(cache%atomic_grid_weights_t)
1647 CALL torch_tensor_release(cache%atomic_grid_sizes_t)
1648 CALL torch_tensor_release(cache%coarse_0_atomic_coords_t)
1649 CALL torch_tensor_release(cache%atomic_grid_size_bound_shape_t)
1650 CALL torch_tensor_release(cache%local_feature_indices_t)
1651 CALL torch_dict_release(cache%static_inputs)
1652 cache%static_tensors_active = .false.
1653 END IF
1654
1655 IF (cache%chunk_static_tensors_active) THEN
1656 CALL torch_tensor_release(cache%chunk_grid_coords_t)
1657 CALL torch_tensor_release(cache%chunk_grid_weights_t)
1658 CALL torch_tensor_release(cache%chunk_atomic_grid_weights_t)
1659 CALL torch_tensor_release(cache%chunk_atomic_grid_sizes_t)
1660 CALL torch_tensor_release(cache%chunk_coarse_0_atomic_coords_t)
1661 CALL torch_tensor_release(cache%chunk_atomic_grid_size_bound_shape_t)
1662 CALL torch_tensor_release(cache%chunk_feature_indices_t)
1663 CALL torch_dict_release(cache%chunk_static_inputs)
1664 cache%chunk_static_tensors_active = .false.
1665 END IF
1666
1667 IF (ALLOCATED(cache%chunk_feature_counts)) DEALLOCATE (cache%chunk_feature_counts)
1668 IF (ALLOCATED(cache%chunk_feature_displs)) DEALLOCATE (cache%chunk_feature_displs)
1669 IF (ALLOCATED(cache%chunk_grad_counts)) DEALLOCATE (cache%chunk_grad_counts)
1670 IF (ALLOCATED(cache%chunk_grad_displs)) DEALLOCATE (cache%chunk_grad_displs)
1671 IF (ALLOCATED(cache%chunk_return_positions)) DEALLOCATE (cache%chunk_return_positions)
1672 IF (ALLOCATED(cache%route_chunk_offsets)) DEALLOCATE (cache%route_chunk_offsets)
1673 IF (ALLOCATED(cache%route_chunk_rows)) DEALLOCATE (cache%route_chunk_rows)
1674 IF (ALLOCATED(cache%route_point_recv_counts)) DEALLOCATE (cache%route_point_recv_counts)
1675 IF (ALLOCATED(cache%route_point_recv_displs)) DEALLOCATE (cache%route_point_recv_displs)
1676 IF (ALLOCATED(cache%route_point_send_counts)) DEALLOCATE (cache%route_point_send_counts)
1677 IF (ALLOCATED(cache%route_point_send_displs)) DEALLOCATE (cache%route_point_send_displs)
1678 IF (ALLOCATED(cache%route_row_offsets)) DEALLOCATE (cache%route_row_offsets)
1679 IF (ALLOCATED(cache%route_row_positions)) DEALLOCATE (cache%route_row_positions)
1680 IF (ALLOCATED(cache%route_send_local_rows)) DEALLOCATE (cache%route_send_local_rows)
1681 IF (ALLOCATED(cache%dynamic_counts)) DEALLOCATE (cache%dynamic_counts)
1682 IF (ALLOCATED(cache%dynamic_displs)) DEALLOCATE (cache%dynamic_displs)
1683 IF (ALLOCATED(cache%feature_counts)) DEALLOCATE (cache%feature_counts)
1684 IF (ALLOCATED(cache%feature_displs)) DEALLOCATE (cache%feature_displs)
1685 IF (ALLOCATED(cache%feature_source_points)) DEALLOCATE (cache%feature_source_points)
1686 IF (ALLOCATED(cache%global_to_feature)) DEALLOCATE (cache%global_to_feature)
1687 IF (ALLOCATED(cache%feature_index)) DEALLOCATE (cache%feature_index)
1688 IF (ALLOCATED(cache%atomic_grid_sizes)) DEALLOCATE (cache%atomic_grid_sizes)
1689 IF (ALLOCATED(cache%chunk_atomic_grid_sizes)) DEALLOCATE (cache%chunk_atomic_grid_sizes)
1690 IF (ALLOCATED(cache%chunk_feature_indices)) DEALLOCATE (cache%chunk_feature_indices)
1691 IF (ALLOCATED(cache%local_feature_counts)) DEALLOCATE (cache%local_feature_counts)
1692 IF (ALLOCATED(cache%local_feature_indices)) DEALLOCATE (cache%local_feature_indices)
1693 IF (ALLOCATED(cache%local_feature_offsets)) DEALLOCATE (cache%local_feature_offsets)
1694 IF (ALLOCATED(cache%local_feature_points)) DEALLOCATE (cache%local_feature_points)
1695 IF (ALLOCATED(cache%local_feature_rows)) DEALLOCATE (cache%local_feature_rows)
1696 IF (ALLOCATED(cache%atomic_grid_size_bound_shape)) THEN
1697 DEALLOCATE (cache%atomic_grid_size_bound_shape)
1698 END IF
1699 IF (ALLOCATED(cache%chunk_atomic_grid_size_bound_shape)) THEN
1700 DEALLOCATE (cache%chunk_atomic_grid_size_bound_shape)
1701 END IF
1702 IF (ALLOCATED(cache%atomic_grid_weights)) DEALLOCATE (cache%atomic_grid_weights)
1703 IF (ALLOCATED(cache%chunk_atomic_grid_weights)) DEALLOCATE (cache%chunk_atomic_grid_weights)
1704 IF (ALLOCATED(cache%chunk_grid_weights)) DEALLOCATE (cache%chunk_grid_weights)
1705 IF (ALLOCATED(cache%grid_weights)) DEALLOCATE (cache%grid_weights)
1706 IF (ALLOCATED(cache%atom_coords)) DEALLOCATE (cache%atom_coords)
1707 IF (ALLOCATED(cache%chunk_coarse_0_atomic_coords)) THEN
1708 DEALLOCATE (cache%chunk_coarse_0_atomic_coords)
1709 END IF
1710 IF (ALLOCATED(cache%coarse_0_atomic_coords)) DEALLOCATE (cache%coarse_0_atomic_coords)
1711 IF (ALLOCATED(cache%chunk_grid_coords)) DEALLOCATE (cache%chunk_grid_coords)
1712 IF (ALLOCATED(cache%grid_coords)) DEALLOCATE (cache%grid_coords)
1713
1714 cache%chunk_atom_begin = 1
1715 cache%chunk_atom_end = 0
1716 cache%chunk_feature_begin = 1
1717 cache%chunk_feature_count = 0
1718 cache%chunk_natom = 0
1719 cache%natom = 0
1720 cache%nflat = 0
1721 cache%nflat_local = 0
1722 cache%npoint = 0
1723 cache%nproc = 0
1724 cache%atom_partition = skala_gpw_atom_partition_hard
1725 cache%bo = 0
1726 cache%bounds = 0
1727 cache%npts = 0
1728 cache%dvol = 0.0_dp
1729 cache%weight_sum = 0.0_dp
1730 cache%weight_sumsq = 0.0_dp
1731 cache%cell_hmat = 0.0_dp
1732 cache%dh = 0.0_dp
1733 cache%active = .false.
1734 cache%has_weights = .false.
1735 cache%chunk_dynamic_tensors_active = .false.
1736 cache%chunk_dynamic_input_views_active = .false.
1737 cache%chunk_inputs_active = .false.
1738 cache%chunk_inputs_use_collapsed_rks = .false.
1739 cache%chunk_static_tensors_active = .false.
1740 cache%dynamic_tensors_active = .false.
1741 cache%inputs_active = .false.
1742 cache%static_tensors_active = .false.
1743
1744 END SUBROUTINE release_layout_cache
1745
1746! **************************************************************************************************
1747!> \brief Release Torch objects and backing arrays owned by a feature bundle.
1748!> \param features ...
1749! **************************************************************************************************
1750 SUBROUTINE skala_gpw_feature_release(features)
1751 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1752
1753 IF (features%active) THEN
1754 IF (features%owns_dynamic_tensors) THEN
1755 IF (features%uses_collapsed_rks_dynamic) THEN
1756 CALL torch_tensor_release(features%density_input_t)
1757 CALL torch_tensor_release(features%grad_input_t)
1758 CALL torch_tensor_release(features%kin_input_t)
1759 END IF
1760 CALL torch_tensor_release(features%density_t)
1761 CALL torch_tensor_release(features%grad_t)
1762 CALL torch_tensor_release(features%kin_t)
1763 END IF
1764 IF (features%owns_static_tensors) THEN
1765 CALL torch_tensor_release(features%grid_coords_t)
1766 CALL torch_tensor_release(features%grid_weights_t)
1767 CALL torch_tensor_release(features%atomic_grid_weights_t)
1768 CALL torch_tensor_release(features%atomic_grid_sizes_t)
1769 CALL torch_tensor_release(features%atomic_grid_size_bound_shape_t)
1770 END IF
1771 IF (features%owns_grid_coordinate_tensor) THEN
1772 CALL torch_tensor_release(features%grid_coords_t)
1773 END IF
1774 IF (features%owns_weight_tensors) THEN
1775 CALL torch_tensor_release(features%grid_weights_t)
1776 CALL torch_tensor_release(features%atomic_grid_weights_t)
1777 END IF
1778 IF (features%owns_static_tensors .OR. features%owns_coordinate_tensor) THEN
1779 CALL torch_tensor_release(features%coarse_0_atomic_coords_t)
1780 END IF
1781 IF (features%owns_inputs) CALL torch_dict_release(features%inputs)
1782 features%active = .false.
1783 features%owns_coordinate_tensor = .false.
1784 features%owns_grid_coordinate_tensor = .false.
1785 features%owns_weight_tensors = .false.
1786 features%owns_dynamic_tensors = .true.
1787 features%owns_inputs = .true.
1788 features%owns_static_tensors = .true.
1789 features%uses_atom_chunk_routing = .false.
1790 features%uses_atom_chunks = .false.
1791 features%uses_collapsed_rks_dynamic = .false.
1792 END IF
1793
1794 IF (ALLOCATED(features%chunk_density)) DEALLOCATE (features%chunk_density)
1795 IF (ALLOCATED(features%chunk_grad)) DEALLOCATE (features%chunk_grad)
1796 IF (ALLOCATED(features%chunk_kin)) DEALLOCATE (features%chunk_kin)
1797 IF (ALLOCATED(features%density)) DEALLOCATE (features%density)
1798 IF (ALLOCATED(features%grad)) DEALLOCATE (features%grad)
1799 IF (ALLOCATED(features%kin)) DEALLOCATE (features%kin)
1800 IF (ALLOCATED(features%chunk_grad_counts)) DEALLOCATE (features%chunk_grad_counts)
1801 IF (ALLOCATED(features%chunk_grad_displs)) DEALLOCATE (features%chunk_grad_displs)
1802 NULLIFY (features%route_point_recv_counts, features%route_point_recv_displs, &
1803 features%route_point_send_counts, features%route_point_send_displs, &
1804 features%route_chunk_offsets, features%route_chunk_rows, &
1805 features%route_row_offsets, features%route_row_positions, &
1806 features%route_send_local_rows)
1807 IF (ALLOCATED(features%feature_index)) DEALLOCATE (features%feature_index)
1808 IF (ALLOCATED(features%local_feature_counts)) DEALLOCATE (features%local_feature_counts)
1809 IF (ALLOCATED(features%local_feature_offsets)) DEALLOCATE (features%local_feature_offsets)
1810 IF (ALLOCATED(features%local_feature_rows)) DEALLOCATE (features%local_feature_rows)
1811 IF (ALLOCATED(features%grid_coords)) DEALLOCATE (features%grid_coords)
1812 IF (ALLOCATED(features%grid_weights)) DEALLOCATE (features%grid_weights)
1813 IF (ALLOCATED(features%atomic_grid_weights)) DEALLOCATE (features%atomic_grid_weights)
1814 IF (ALLOCATED(features%atomic_grid_sizes)) DEALLOCATE (features%atomic_grid_sizes)
1815 IF (ALLOCATED(features%coarse_0_atomic_coords)) DEALLOCATE (features%coarse_0_atomic_coords)
1816 IF (ALLOCATED(features%atomic_grid_size_bound_shape)) THEN
1817 DEALLOCATE (features%atomic_grid_size_bound_shape)
1818 END IF
1819 features%chunk_feature_count = 0
1820 features%nflat = 0
1821 features%nflat_local = 0
1822 features%atom_partition = skala_gpw_atom_partition_hard
1823 features%uses_atom_chunk_routing = .false.
1824 features%uses_collapsed_rks_dynamic = .false.
1825
1826 END SUBROUTINE skala_gpw_feature_release
1827
1828! **************************************************************************************************
1829!> \brief Return how many atom-contiguous subchunks the cached rank chunk needs.
1830!> \param max_rows ...
1831!> \return ...
1832! **************************************************************************************************
1833 FUNCTION skala_gpw_atom_subchunk_count(max_rows) RESULT(nsubchunks)
1834 INTEGER, INTENT(IN) :: max_rows
1835 INTEGER :: nsubchunks
1836
1837 INTEGER :: atom_rows, iatom, rows
1838
1839 nsubchunks = 0
1840 IF (.NOT. cached_layout%active) RETURN
1841 IF (cached_layout%chunk_natom <= 0) RETURN
1842 IF (max_rows <= 0) THEN
1843 nsubchunks = 1
1844 RETURN
1845 END IF
1846
1847 rows = 0
1848 DO iatom = 1, cached_layout%chunk_natom
1849 atom_rows = int(cached_layout%chunk_atomic_grid_sizes(iatom))
1850 IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1851 nsubchunks = nsubchunks + 1
1852 rows = 0
1853 END IF
1854 rows = rows + atom_rows
1855 END DO
1856 IF (rows > 0) nsubchunks = nsubchunks + 1
1857 nsubchunks = max(1, nsubchunks)
1858
1859 END FUNCTION skala_gpw_atom_subchunk_count
1860
1861! **************************************************************************************************
1862!> \brief Build atom and row bounds for all atom-contiguous subchunks in one pass.
1863!> \param max_rows ...
1864!> \param atom_begin ...
1865!> \param atom_count ...
1866!> \param row_begin ...
1867!> \param row_count ...
1868! **************************************************************************************************
1869 SUBROUTINE skala_gpw_atom_subchunk_layout(max_rows, atom_begin, atom_count, row_begin, &
1870 row_count)
1871 INTEGER, INTENT(IN) :: max_rows
1872 INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: atom_begin, atom_count, row_begin, &
1873 row_count
1874
1875 INTEGER :: atom_begin_tmp, atom_count_tmp, atom_rows, iatom, insert_at, nsubchunks, &
1876 row_begin_tmp, row_count_tmp, rows, subchunk
1877 INTEGER(KIND=int_8) :: padded_rows_tmp
1878 INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: padded_rows
1879
1880 nsubchunks = skala_gpw_atom_subchunk_count(max_rows)
1881 ALLOCATE (atom_begin(nsubchunks), atom_count(nsubchunks), row_begin(nsubchunks), &
1882 row_count(nsubchunks))
1883 IF (nsubchunks == 0) RETURN
1884
1885 cpassert(cached_layout%active)
1886 cpassert(cached_layout%chunk_natom > 0)
1887
1888 atom_begin = 0
1889 atom_count = 0
1890 row_begin = 0
1891 row_count = 0
1892
1893 IF (max_rows <= 0) THEN
1894 atom_begin(1) = 1
1895 atom_count(1) = cached_layout%chunk_natom
1896 row_begin(1) = 1
1897 row_count(1) = cached_layout%chunk_feature_count
1898 RETURN
1899 END IF
1900
1901 subchunk = 1
1902 atom_begin(subchunk) = 1
1903 row_begin(subchunk) = 1
1904 rows = 0
1905 DO iatom = 1, cached_layout%chunk_natom
1906 atom_rows = int(cached_layout%chunk_atomic_grid_sizes(iatom))
1907 IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1908 atom_count(subchunk) = iatom - atom_begin(subchunk)
1909 row_count(subchunk) = rows
1910 subchunk = subchunk + 1
1911 atom_begin(subchunk) = iatom
1912 row_begin(subchunk) = row_begin(subchunk - 1) + row_count(subchunk - 1)
1913 rows = 0
1914 END IF
1915 rows = rows + atom_rows
1916 END DO
1917 atom_count(subchunk) = cached_layout%chunk_natom - atom_begin(subchunk) + 1
1918 row_count(subchunk) = rows
1919
1920 cpassert(subchunk == nsubchunks)
1921
1922 ALLOCATE (padded_rows(nsubchunks))
1923 DO subchunk = 1, nsubchunks
1924 iatom = atom_begin(subchunk) + atom_count(subchunk) - 1
1925 padded_rows(subchunk) = int(atom_count(subchunk), kind=int_8)* &
1926 maxval(cached_layout%chunk_atomic_grid_sizes( &
1927 atom_begin(subchunk):iatom))
1928 END DO
1929
1930 ! Evaluate the largest padded block first so Torch can reuse its CUDA allocations.
1931 DO subchunk = 2, nsubchunks
1932 padded_rows_tmp = padded_rows(subchunk)
1933 atom_begin_tmp = atom_begin(subchunk)
1934 atom_count_tmp = atom_count(subchunk)
1935 row_begin_tmp = row_begin(subchunk)
1936 row_count_tmp = row_count(subchunk)
1937 insert_at = subchunk
1938 DO WHILE (insert_at > 1)
1939 IF (padded_rows(insert_at - 1) >= padded_rows_tmp) EXIT
1940 padded_rows(insert_at) = padded_rows(insert_at - 1)
1941 atom_begin(insert_at) = atom_begin(insert_at - 1)
1942 atom_count(insert_at) = atom_count(insert_at - 1)
1943 row_begin(insert_at) = row_begin(insert_at - 1)
1944 row_count(insert_at) = row_count(insert_at - 1)
1945 insert_at = insert_at - 1
1946 END DO
1947 padded_rows(insert_at) = padded_rows_tmp
1948 atom_begin(insert_at) = atom_begin_tmp
1949 atom_count(insert_at) = atom_count_tmp
1950 row_begin(insert_at) = row_begin_tmp
1951 row_count(insert_at) = row_count_tmp
1952 END DO
1953 DEALLOCATE (padded_rows)
1954
1955 END SUBROUTINE skala_gpw_atom_subchunk_layout
1956
1957! **************************************************************************************************
1958!> \brief Build an atom-contiguous subchunk feature bundle from a rank-local atom chunk.
1959!> \param parent ...
1960!> \param features ...
1961!> \param subchunk_index ...
1962!> \param max_rows ...
1963!> \param requires_grad ...
1964! **************************************************************************************************
1965 SUBROUTINE skala_gpw_feature_build_atom_subchunk(parent, features, subchunk_index, &
1966 max_rows, requires_grad)
1967 TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1968 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1969 INTEGER, INTENT(IN) :: subchunk_index, max_rows
1970 LOGICAL, INTENT(IN) :: requires_grad
1971
1972 INTEGER :: atom_begin, atom_end, row_begin, row_end
1973
1974 cpassert(parent%uses_atom_chunks)
1975 CALL atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
1976 row_begin, row_end)
1977 CALL skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, &
1978 atom_end - atom_begin + 1, row_begin, &
1979 row_end - row_begin + 1, requires_grad)
1980
1981 END SUBROUTINE skala_gpw_feature_build_atom_subchunk
1982
1983! **************************************************************************************************
1984!> \brief Build an atom-contiguous subchunk feature bundle from precomputed bounds.
1985!> \param parent ...
1986!> \param features ...
1987!> \param atom_begin ...
1988!> \param atom_count ...
1989!> \param row_begin ...
1990!> \param row_count ...
1991!> \param requires_grad ...
1992! **************************************************************************************************
1993 SUBROUTINE skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, &
1994 atom_count, row_begin, row_count, &
1995 requires_grad)
1996 TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1997 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1998 INTEGER, INTENT(IN) :: atom_begin, atom_count, row_begin, &
1999 row_count
2000 LOGICAL, INTENT(IN) :: requires_grad
2001
2002 INTEGER :: atom_end, max_grid_size, row_end
2003
2004 CALL skala_gpw_feature_release(features)
2005 cpassert(parent%uses_atom_chunks)
2006 atom_end = atom_begin + atom_count - 1
2007 row_end = row_begin + row_count - 1
2008 cpassert(atom_count > 0)
2009 cpassert(row_count > 0)
2010 mark_used(requires_grad)
2011 max_grid_size = maxval(int(cached_layout%chunk_atomic_grid_sizes(atom_begin:atom_end)))
2012
2013 ALLOCATE (features%atomic_grid_size_bound_shape(0, max_grid_size))
2014 features%atomic_grid_size_bound_shape = 0_int_8
2015
2016 features%chunk_feature_count = row_count
2017 features%nflat = parent%nflat
2018 features%nflat_local = parent%nflat_local
2019 features%grid_weight_sum = sum(cached_layout%chunk_grid_weights(row_begin:row_end))
2020 features%uses_atom_chunks = .true.
2021 features%uses_atom_chunk_routing = parent%uses_atom_chunk_routing
2022 CALL add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
2023 row_count)
2024 features%active = .true.
2025
2026 END SUBROUTINE skala_gpw_feature_build_atom_subchunk_bounds
2027
2028! **************************************************************************************************
2029!> \brief Return atom and row bounds for an atom-contiguous rank-local subchunk.
2030!> \param subchunk_index ...
2031!> \param max_rows ...
2032!> \param atom_begin ...
2033!> \param atom_end ...
2034!> \param row_begin ...
2035!> \param row_end ...
2036! **************************************************************************************************
2037 SUBROUTINE atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
2038 row_begin, row_end)
2039 INTEGER, INTENT(IN) :: subchunk_index, max_rows
2040 INTEGER, INTENT(OUT) :: atom_begin, atom_end, row_begin, row_end
2041
2042 INTEGER :: atom_rows, current_subchunk, iatom, &
2043 row_cursor, rows
2044
2045 cpassert(subchunk_index > 0)
2046 cpassert(max_rows > 0)
2047 cpassert(cached_layout%chunk_natom > 0)
2048
2049 atom_begin = 1
2050 atom_end = 0
2051 row_begin = 1
2052 row_end = 0
2053 current_subchunk = 1
2054 row_cursor = 1
2055 rows = 0
2056 DO iatom = 1, cached_layout%chunk_natom
2057 atom_rows = int(cached_layout%chunk_atomic_grid_sizes(iatom))
2058 IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
2059 IF (current_subchunk == subchunk_index) THEN
2060 atom_end = iatom - 1
2061 row_end = row_cursor - 1
2062 RETURN
2063 END IF
2064 current_subchunk = current_subchunk + 1
2065 atom_begin = iatom
2066 row_begin = row_cursor
2067 rows = 0
2068 END IF
2069 rows = rows + atom_rows
2070 row_cursor = row_cursor + atom_rows
2071 END DO
2072
2073 IF (current_subchunk == subchunk_index) THEN
2074 atom_end = cached_layout%chunk_natom
2075 row_end = row_cursor - 1
2076 RETURN
2077 END IF
2078
2079 cpabort("Requested native SKALA atom subchunk does not exist.")
2080
2081 END SUBROUTINE atom_subchunk_bounds
2082
2083! **************************************************************************************************
2084!> \brief Insert a subchunk into a Torch dictionary using static views of the cached chunk tensors.
2085!> \param parent ...
2086!> \param features ...
2087!> \param atom_begin ...
2088!> \param atom_count ...
2089!> \param row_begin ...
2090!> \param row_count ...
2091! **************************************************************************************************
2092 SUBROUTINE add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
2093 row_count)
2094 TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
2095 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2096 INTEGER, INTENT(IN) :: atom_begin, atom_count, row_begin, &
2097 row_count
2098
2099 cpassert(cached_layout%chunk_static_tensors_active)
2100 cpassert(parent%active)
2101 cpassert(ALLOCATED(features%atomic_grid_size_bound_shape))
2102
2103 features%owns_coordinate_tensor = .false.
2104 features%owns_dynamic_tensors = .true.
2105 features%owns_inputs = .true.
2106 features%owns_static_tensors = .true.
2107 features%uses_collapsed_rks_dynamic = parent%uses_collapsed_rks_dynamic
2108
2109 CALL torch_tensor_narrow(cached_layout%chunk_grid_coords_t, 0, row_begin - 1, &
2110 row_count, features%grid_coords_t)
2111 CALL torch_tensor_narrow(cached_layout%chunk_grid_weights_t, 0, row_begin - 1, &
2112 row_count, features%grid_weights_t)
2113 CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_weights_t, 0, row_begin - 1, &
2114 row_count, features%atomic_grid_weights_t)
2115 CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_sizes_t, 0, atom_begin - 1, &
2116 atom_count, features%atomic_grid_sizes_t)
2117 CALL torch_tensor_narrow(cached_layout%chunk_coarse_0_atomic_coords_t, 0, &
2118 atom_begin - 1, atom_count, features%coarse_0_atomic_coords_t)
2119 CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
2120 features%atomic_grid_size_bound_shape)
2121 CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .false.)
2122 CALL torch_tensor_narrow(parent%density_t, 1, row_begin - 1, row_count, &
2123 features%density_t)
2124 CALL torch_tensor_narrow(parent%grad_t, 2, row_begin - 1, row_count, features%grad_t)
2125 CALL torch_tensor_narrow(parent%kin_t, 1, row_begin - 1, row_count, features%kin_t)
2126 IF (features%uses_collapsed_rks_dynamic) THEN
2127 CALL torch_tensor_expand_dim(features%density_t, 0, 2, features%density_input_t)
2128 CALL torch_tensor_expand_dim(features%grad_t, 0, 2, features%grad_input_t)
2129 CALL torch_tensor_expand_dim(features%kin_t, 0, 2, features%kin_input_t)
2130 END IF
2131
2132 CALL torch_dict_create(features%inputs)
2133 CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2134 CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2135 CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2136 features%atomic_grid_weights_t)
2137 CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2138 features%atomic_grid_sizes_t)
2139 CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2140 features%atomic_grid_size_bound_shape_t)
2141 IF (features%uses_collapsed_rks_dynamic) THEN
2142 CALL torch_dict_insert(features%inputs, "density", features%density_input_t)
2143 CALL torch_dict_insert(features%inputs, "grad", features%grad_input_t)
2144 CALL torch_dict_insert(features%inputs, "kin", features%kin_input_t)
2145 ELSE
2146 CALL torch_dict_insert(features%inputs, "density", features%density_t)
2147 CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2148 CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2149 END IF
2150 CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2151 features%coarse_0_atomic_coords_t)
2152
2153 END SUBROUTINE add_subchunk_feature_tensors
2154
2155! **************************************************************************************************
2156!> \brief Insert owned subchunk arrays into a Torch dictionary.
2157!> \param features ...
2158!> \param requires_grad ...
2159! **************************************************************************************************
2160 SUBROUTINE add_owned_feature_tensors(features, requires_grad)
2161 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2162 LOGICAL, INTENT(IN) :: requires_grad
2163
2164 cpassert(ALLOCATED(features%chunk_density))
2165 cpassert(ALLOCATED(features%chunk_grad))
2166 cpassert(ALLOCATED(features%chunk_kin))
2167 cpassert(ALLOCATED(features%grid_coords))
2168 cpassert(ALLOCATED(features%grid_weights))
2169 cpassert(ALLOCATED(features%atomic_grid_weights))
2170 cpassert(ALLOCATED(features%atomic_grid_sizes))
2171 cpassert(ALLOCATED(features%atomic_grid_size_bound_shape))
2172 cpassert(ALLOCATED(features%coarse_0_atomic_coords))
2173
2174 features%owns_coordinate_tensor = .false.
2175 features%owns_dynamic_tensors = .true.
2176 features%owns_inputs = .true.
2177 features%owns_static_tensors = .true.
2178
2179 CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
2180 CALL torch_tensor_to_device_leaf(features%grid_coords_t, .false.)
2181 CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
2182 CALL torch_tensor_to_device_leaf(features%grid_weights_t, .false.)
2183 CALL torch_tensor_from_array(features%atomic_grid_weights_t, features%atomic_grid_weights)
2184 CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .false.)
2185 CALL torch_tensor_from_array(features%atomic_grid_sizes_t, features%atomic_grid_sizes)
2186 CALL torch_tensor_to_device_leaf(features%atomic_grid_sizes_t, .false.)
2187 CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
2188 features%coarse_0_atomic_coords)
2189 CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .false.)
2190 CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
2191 features%atomic_grid_size_bound_shape)
2192 CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .false.)
2193 CALL torch_tensor_from_array(features%density_t, features%chunk_density)
2194 CALL torch_tensor_to_device_leaf(features%density_t, requires_grad)
2195 CALL torch_tensor_from_array(features%grad_t, features%chunk_grad)
2196 CALL torch_tensor_to_device_leaf(features%grad_t, requires_grad)
2197 CALL torch_tensor_from_array(features%kin_t, features%chunk_kin)
2198 CALL torch_tensor_to_device_leaf(features%kin_t, requires_grad)
2199
2200 CALL torch_dict_create(features%inputs)
2201 CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2202 CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2203 CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2204 features%atomic_grid_weights_t)
2205 CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2206 features%atomic_grid_sizes_t)
2207 CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2208 features%atomic_grid_size_bound_shape_t)
2209 CALL torch_dict_insert(features%inputs, "density", features%density_t)
2210 CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2211 CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2212 CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2213 features%coarse_0_atomic_coords_t)
2214
2215 END SUBROUTINE add_owned_feature_tensors
2216
2217! **************************************************************************************************
2218!> \brief Insert all SKALA feature tensors into the Torch dictionary.
2219!> \param features ...
2220!> \param requires_grad ...
2221!> \param requires_coordinate_grad ...
2222!> \param requires_stress_grad ...
2223!> \param use_atom_chunks ...
2224!> \param requires_weight_grad ...
2225! **************************************************************************************************
2226 SUBROUTINE add_feature_tensors(features, requires_grad, requires_coordinate_grad, &
2227 requires_stress_grad, use_atom_chunks, requires_weight_grad)
2228 TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
2229 LOGICAL, INTENT(IN) :: requires_grad, requires_coordinate_grad, &
2230 requires_stress_grad, use_atom_chunks
2231 LOGICAL, INTENT(IN), OPTIONAL :: requires_weight_grad
2232
2233 LOGICAL :: my_requires_weight_grad
2234
2235 my_requires_weight_grad = .false.
2236 IF (PRESENT(requires_weight_grad)) my_requires_weight_grad = requires_weight_grad
2237
2238 features%owns_static_tensors = .false.
2239 features%owns_coordinate_tensor = .false.
2240 features%owns_grid_coordinate_tensor = .false.
2241 features%owns_weight_tensors = .false.
2242 features%owns_dynamic_tensors = .false.
2243 features%owns_inputs = .true.
2244 IF (use_atom_chunks) THEN
2245 cpassert(.NOT. requires_coordinate_grad)
2246 cpassert(.NOT. requires_stress_grad)
2247 cpassert(.NOT. my_requires_weight_grad)
2248 cpassert(cached_layout%chunk_static_tensors_active)
2249 features%grid_coords_t = cached_layout%chunk_grid_coords_t
2250 features%grid_weights_t = cached_layout%chunk_grid_weights_t
2251 features%atomic_grid_weights_t = cached_layout%chunk_atomic_grid_weights_t
2252 features%atomic_grid_sizes_t = cached_layout%chunk_atomic_grid_sizes_t
2253 features%atomic_grid_size_bound_shape_t = &
2254 cached_layout%chunk_atomic_grid_size_bound_shape_t
2255 features%local_feature_indices_t = cached_layout%chunk_feature_indices_t
2256
2257 IF (cached_layout%chunk_inputs_active .AND. &
2258 (cached_layout%chunk_inputs_use_collapsed_rks .NEQV. &
2259 features%uses_collapsed_rks_dynamic)) THEN
2260 CALL torch_dict_release(cached_layout%chunk_inputs)
2261 cached_layout%chunk_inputs_active = .false.
2262 END IF
2263 IF (.NOT. features%uses_collapsed_rks_dynamic .AND. &
2264 cached_layout%chunk_dynamic_input_views_active) THEN
2265 CALL torch_tensor_release(cached_layout%chunk_density_input_t)
2266 CALL torch_tensor_release(cached_layout%chunk_grad_input_t)
2267 CALL torch_tensor_release(cached_layout%chunk_kin_input_t)
2268 cached_layout%chunk_dynamic_input_views_active = .false.
2269 END IF
2270
2271 CALL torch_tensor_reset_from_array(cached_layout%chunk_density_t, &
2272 features%chunk_density, requires_grad=requires_grad)
2273 features%density_t = cached_layout%chunk_density_t
2274 CALL torch_tensor_reset_from_array(cached_layout%chunk_grad_t, features%chunk_grad, &
2275 requires_grad=requires_grad)
2276 features%grad_t = cached_layout%chunk_grad_t
2277 CALL torch_tensor_reset_from_array(cached_layout%chunk_kin_t, features%chunk_kin, &
2278 requires_grad=requires_grad)
2279 features%kin_t = cached_layout%chunk_kin_t
2280 cached_layout%chunk_dynamic_tensors_active = .true.
2281
2282 IF (features%uses_collapsed_rks_dynamic .AND. &
2283 .NOT. cached_layout%chunk_dynamic_input_views_active) THEN
2284 CALL torch_tensor_expand_dim(cached_layout%chunk_density_t, 0, 2, &
2285 cached_layout%chunk_density_input_t)
2286 CALL torch_tensor_expand_dim(cached_layout%chunk_grad_t, 0, 2, &
2287 cached_layout%chunk_grad_input_t)
2288 CALL torch_tensor_expand_dim(cached_layout%chunk_kin_t, 0, 2, &
2289 cached_layout%chunk_kin_input_t)
2290 cached_layout%chunk_dynamic_input_views_active = .true.
2291 END IF
2292 IF (features%uses_collapsed_rks_dynamic) THEN
2293 features%density_input_t = cached_layout%chunk_density_input_t
2294 features%grad_input_t = cached_layout%chunk_grad_input_t
2295 features%kin_input_t = cached_layout%chunk_kin_input_t
2296 END IF
2297
2298 IF (.NOT. cached_layout%chunk_inputs_active) THEN
2299 CALL torch_dict_clone(cached_layout%chunk_static_inputs, cached_layout%chunk_inputs)
2300 IF (features%uses_collapsed_rks_dynamic) THEN
2301 CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2302 features%density_input_t)
2303 CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2304 features%grad_input_t)
2305 CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2306 features%kin_input_t)
2307 ELSE
2308 CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2309 cached_layout%chunk_density_t)
2310 CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2311 cached_layout%chunk_grad_t)
2312 CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2313 cached_layout%chunk_kin_t)
2314 END IF
2315 CALL torch_dict_insert(cached_layout%chunk_inputs, "coarse_0_atomic_coords", &
2316 cached_layout%chunk_coarse_0_atomic_coords_t)
2317 cached_layout%chunk_inputs_use_collapsed_rks = features%uses_collapsed_rks_dynamic
2318 cached_layout%chunk_inputs_active = .true.
2319 END IF
2320 features%inputs = cached_layout%chunk_inputs
2321 features%owns_inputs = .false.
2322 features%coarse_0_atomic_coords_t = cached_layout%chunk_coarse_0_atomic_coords_t
2323 ELSE
2324 cpassert(cached_layout%static_tensors_active)
2325 IF (.NOT. requires_stress_grad .AND. .NOT. my_requires_weight_grad) THEN
2326 features%grid_coords_t = cached_layout%grid_coords_t
2327 features%grid_weights_t = cached_layout%grid_weights_t
2328 features%atomic_grid_weights_t = cached_layout%atomic_grid_weights_t
2329 END IF
2330 features%atomic_grid_sizes_t = cached_layout%atomic_grid_sizes_t
2331 features%atomic_grid_size_bound_shape_t = cached_layout%atomic_grid_size_bound_shape_t
2332 features%local_feature_indices_t = cached_layout%local_feature_indices_t
2333
2334 CALL torch_tensor_reset_from_array(cached_layout%density_t, features%density, &
2335 requires_grad=requires_grad)
2336 features%density_t = cached_layout%density_t
2337 CALL torch_tensor_reset_from_array(cached_layout%grad_t, features%grad, &
2338 requires_grad=requires_grad)
2339 features%grad_t = cached_layout%grad_t
2340 CALL torch_tensor_reset_from_array(cached_layout%kin_t, features%kin, &
2341 requires_grad=requires_grad)
2342 features%kin_t = cached_layout%kin_t
2343 cached_layout%dynamic_tensors_active = .true.
2344
2345 IF (requires_coordinate_grad .OR. requires_stress_grad .OR. my_requires_weight_grad) THEN
2346 IF (requires_stress_grad .OR. my_requires_weight_grad) THEN
2347 CALL torch_dict_create(features%inputs)
2348 IF (requires_stress_grad) THEN
2349 CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
2350 CALL torch_tensor_to_device_leaf(features%grid_coords_t, .true.)
2351 CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2352 features%owns_grid_coordinate_tensor = .true.
2353 ELSE
2354 features%grid_coords_t = cached_layout%grid_coords_t
2355 CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2356 END IF
2357 CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
2358 CALL torch_tensor_to_device_leaf(features%grid_weights_t, .true.)
2359 CALL torch_tensor_from_array(features%atomic_grid_weights_t, &
2360 features%atomic_grid_weights)
2361 CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .true.)
2362 CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2363 CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2364 features%atomic_grid_weights_t)
2365 CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2366 features%atomic_grid_sizes_t)
2367 CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2368 features%atomic_grid_size_bound_shape_t)
2369 features%owns_weight_tensors = .true.
2370 ELSE
2371 CALL torch_dict_clone(cached_layout%static_inputs, features%inputs)
2372 END IF
2373 CALL torch_dict_insert(features%inputs, "density", features%density_t)
2374 CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2375 CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2376 ELSE
2377 IF (.NOT. cached_layout%inputs_active) THEN
2378 CALL torch_dict_clone(cached_layout%static_inputs, cached_layout%inputs)
2379 CALL torch_dict_insert(cached_layout%inputs, "density", cached_layout%density_t)
2380 CALL torch_dict_insert(cached_layout%inputs, "grad", cached_layout%grad_t)
2381 CALL torch_dict_insert(cached_layout%inputs, "kin", cached_layout%kin_t)
2382 CALL torch_dict_insert(cached_layout%inputs, "coarse_0_atomic_coords", &
2383 cached_layout%coarse_0_atomic_coords_t)
2384 cached_layout%inputs_active = .true.
2385 END IF
2386 features%inputs = cached_layout%inputs
2387 features%owns_inputs = .false.
2388 features%coarse_0_atomic_coords_t = cached_layout%coarse_0_atomic_coords_t
2389 END IF
2390 END IF
2391
2392 IF (requires_coordinate_grad .OR. requires_stress_grad) THEN
2393 cpassert(.NOT. use_atom_chunks)
2394 CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
2395 features%coarse_0_atomic_coords)
2396 CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .true.)
2397 CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2398 features%coarse_0_atomic_coords_t)
2399 features%owns_coordinate_tensor = .true.
2400 END IF
2401
2402 END SUBROUTINE add_feature_tensors
2403
2404! **************************************************************************************************
2405!> \brief Return the Cartesian coordinate of a regular GPW grid point.
2406!> \param pw_grid ...
2407!> \param index ...
2408!> \return ...
2409! **************************************************************************************************
2410 FUNCTION grid_coordinate(pw_grid, index) RESULT(coord)
2411 TYPE(pw_grid_type), POINTER :: pw_grid
2412 INTEGER, DIMENSION(3), INTENT(IN) :: index
2413 REAL(kind=dp), DIMENSION(3) :: coord
2414
2415 INTEGER, DIMENSION(3) :: relative_index
2416
2417 relative_index = index - pw_grid%bounds(1, :)
2418 coord = real(relative_index(1), kind=dp)*pw_grid%dh(:, 1) + &
2419 REAL(relative_index(2), kind=dp)*pw_grid%dh(:, 2) + &
2420 REAL(relative_index(3), kind=dp)*pw_grid%dh(:, 3)
2421
2422 END FUNCTION grid_coordinate
2423
2424! **************************************************************************************************
2425!> \brief Build Becke-like smooth atom weights for one native-grid point.
2426!> \param grid_point ...
2427!> \param atom_coords ...
2428!> \param cell ...
2429!> \param weights ...
2430!> \param partition_atom_coords ...
2431!> \param distances ...
2432!> \param pair_distances ...
2433! **************************************************************************************************
2434 SUBROUTINE smooth_atom_partition(grid_point, atom_coords, cell, weights, partition_atom_coords, &
2435 distances, pair_distances)
2436 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
2437 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2438 TYPE(cell_type), POINTER :: cell
2439 REAL(kind=dp), DIMENSION(:), INTENT(OUT) :: weights
2440 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT) :: partition_atom_coords
2441 REAL(kind=dp), DIMENSION(:), INTENT(OUT) :: distances
2442 REAL(kind=dp), DIMENSION(:, :), INTENT(IN), &
2443 OPTIONAL :: pair_distances
2444
2445 INTEGER :: iatom, jatom, natom
2446 REAL(kind=dp) :: mu, rab, rsum, switch, total
2447 REAL(kind=dp), DIMENSION(3) :: rij
2448
2449 natom = SIZE(atom_coords, 2)
2450 cpassert(SIZE(weights) == natom)
2451 cpassert(SIZE(partition_atom_coords, 1) == 3)
2452 cpassert(SIZE(partition_atom_coords, 2) == natom)
2453 cpassert(SIZE(distances) == natom)
2454
2455 IF (PRESENT(pair_distances)) THEN
2456 DO iatom = 1, natom
2457 rij = grid_point - atom_coords(:, iatom)
2458 distances(iatom) = sqrt(sum(rij**2))
2459 END DO
2460 ELSE
2461 DO iatom = 1, natom
2462 partition_atom_coords(:, iatom) = &
2463 nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2464 rij = grid_point - partition_atom_coords(:, iatom)
2465 distances(iatom) = sqrt(sum(rij**2))
2466 END DO
2467 END IF
2468
2469 weights = 1.0_dp
2470 DO iatom = 1, natom - 1
2471 DO jatom = iatom + 1, natom
2472 IF (PRESENT(pair_distances)) THEN
2473 rab = pair_distances(iatom, jatom)
2474 ELSE
2475 rij = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
2476 rab = sqrt(sum(rij**2))
2477 END IF
2478 IF (rab <= layout_tol) cycle
2479 mu = (distances(iatom) - distances(jatom))/rab
2480 mu = max(-1.0_dp, min(1.0_dp, mu))
2481 switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2482 weights(iatom) = weights(iatom)*switch
2483 weights(jatom) = weights(jatom)*(1.0_dp - switch)
2484 END DO
2485 END DO
2486
2487 total = sum(weights)
2488 IF (total > 0.0_dp) THEN
2489 weights = weights/total
2490 ELSE
2491 rsum = huge(1.0_dp)
2492 jatom = 1
2493 DO iatom = 1, natom
2494 IF (distances(iatom) < rsum) THEN
2495 rsum = distances(iatom)
2496 jatom = iatom
2497 END IF
2498 END DO
2499 weights = 0.0_dp
2500 weights(jatom) = 1.0_dp
2501 END IF
2502
2503 END SUBROUTINE smooth_atom_partition
2504
2505! **************************************************************************************************
2506!> \brief Build screened logarithmic Becke weights for a set of periodic atom images.
2507!> \param image_coords ...
2508!> \param distances ...
2509!> \param log_weights ...
2510! **************************************************************************************************
2511 SUBROUTINE periodic_image_log_weights(image_coords, distances, log_weights)
2512 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: image_coords
2513 REAL(kind=dp), DIMENSION(:), INTENT(IN) :: distances
2514 REAL(kind=dp), DIMENSION(:), INTENT(OUT) :: log_weights
2515
2516 INTEGER :: iimage, image_count, jimage, &
2517 nearest_image
2518 REAL(kind=dp) :: factor, mu, rab, reference_log_weight, &
2519 screen_log_weight
2520 REAL(kind=dp), DIMENSION(3) :: pair
2521
2522 image_count = SIZE(distances)
2523 cpassert(SIZE(image_coords, 1) == 3)
2524 cpassert(SIZE(image_coords, 2) == image_count)
2525 cpassert(SIZE(log_weights) == image_count)
2526
2527 nearest_image = minloc(distances, dim=1)
2528 log_weights = -huge(1.0_dp)
2529 reference_log_weight = 0.0_dp
2530 DO jimage = 1, image_count
2531 IF (jimage == nearest_image) cycle
2532 pair = image_coords(:, nearest_image) - image_coords(:, jimage)
2533 rab = sqrt(sum(pair**2))
2534 IF (rab <= layout_tol) cycle
2535 mu = max(-1.0_dp, min(1.0_dp, &
2536 (distances(nearest_image) - distances(jimage))/rab))
2537 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2538 cpassert(factor > tiny(1.0_dp))
2539 reference_log_weight = reference_log_weight + log(factor)
2540 END DO
2541 log_weights(nearest_image) = reference_log_weight
2542
2543 ! Every omitted factor is at most one. A partial log product is therefore an upper bound
2544 ! on the final unnormalized weight. The normalization contains at least the nearest-image
2545 ! reference weight, so screening below this limit changes the complete partition by no more
2546 ! than one floating-point epsilon in total.
2547 screen_log_weight = reference_log_weight + &
2548 log(epsilon(1.0_dp)/real(image_count, kind=dp))
2549 DO iimage = 1, image_count
2550 IF (iimage == nearest_image) cycle
2551 pair = image_coords(:, iimage) - image_coords(:, nearest_image)
2552 rab = sqrt(sum(pair**2))
2553 IF (rab <= layout_tol) cycle
2554 mu = max(-1.0_dp, min(1.0_dp, &
2555 (distances(iimage) - distances(nearest_image))/rab))
2556 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2557 IF (factor <= tiny(1.0_dp)) cycle
2558 log_weights(iimage) = log(factor)
2559 IF (log_weights(iimage) <= screen_log_weight) THEN
2560 log_weights(iimage) = -huge(1.0_dp)
2561 cycle
2562 END IF
2563
2564 DO jimage = 1, image_count
2565 IF (jimage == iimage .OR. jimage == nearest_image) cycle
2566 pair = image_coords(:, iimage) - image_coords(:, jimage)
2567 rab = sqrt(sum(pair**2))
2568 IF (rab <= layout_tol) cycle
2569 mu = max(-1.0_dp, min(1.0_dp, &
2570 (distances(iimage) - distances(jimage))/rab))
2571 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2572 IF (factor <= tiny(1.0_dp)) THEN
2573 log_weights(iimage) = -huge(1.0_dp)
2574 EXIT
2575 END IF
2576 log_weights(iimage) = log_weights(iimage) + log(factor)
2577 IF (log_weights(iimage) <= screen_log_weight) THEN
2578 log_weights(iimage) = -huge(1.0_dp)
2579 EXIT
2580 END IF
2581 END DO
2582 END DO
2583
2584 END SUBROUTINE periodic_image_log_weights
2585
2586! **************************************************************************************************
2587!> \brief Build screened normalized Becke weights, using direct products whenever they are safe.
2588!> \param image_coords ...
2589!> \param distances ...
2590!> \param normalized_weights ...
2591! **************************************************************************************************
2592 SUBROUTINE periodic_image_normalized_weights(image_coords, distances, normalized_weights)
2593 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: image_coords
2594 REAL(kind=dp), DIMENSION(:), INTENT(IN) :: distances
2595 REAL(kind=dp), DIMENSION(:), INTENT(OUT) :: normalized_weights
2596
2597 INTEGER :: iimage, image_count, jimage, &
2598 nearest_image
2599 LOGICAL :: use_log_fallback
2600 REAL(kind=dp) :: factor, max_log_weight, mu, &
2601 normalization, product, rab, &
2602 reference_weight, screen_factor, &
2603 screen_weight
2604 REAL(kind=dp), DIMENSION(3) :: pair
2605 REAL(kind=dp), DIMENSION(SIZE(distances)) :: log_weights, raw_weights
2606
2607 image_count = SIZE(distances)
2608 cpassert(SIZE(image_coords, 1) == 3)
2609 cpassert(SIZE(image_coords, 2) == image_count)
2610 cpassert(SIZE(normalized_weights) == image_count)
2611 IF (image_count == 1) THEN
2612 normalized_weights(1) = 1.0_dp
2613 RETURN
2614 END IF
2615
2616 nearest_image = minloc(distances, dim=1)
2617 reference_weight = 1.0_dp
2618 use_log_fallback = .false.
2619 DO jimage = 1, image_count
2620 IF (jimage == nearest_image) cycle
2621 pair = image_coords(:, nearest_image) - image_coords(:, jimage)
2622 rab = sqrt(sum(pair**2))
2623 IF (rab <= layout_tol) cycle
2624 mu = max(-1.0_dp, min(1.0_dp, &
2625 (distances(nearest_image) - distances(jimage))/rab))
2626 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2627 cpassert(factor > tiny(1.0_dp))
2628 IF (reference_weight <= tiny(1.0_dp)/factor) THEN
2629 use_log_fallback = .true.
2630 EXIT
2631 END IF
2632 reference_weight = reference_weight*factor
2633 END DO
2634
2635 screen_factor = epsilon(1.0_dp)/real(image_count, kind=dp)
2636 IF (.NOT. use_log_fallback) THEN
2637 use_log_fallback = reference_weight <= tiny(1.0_dp)/screen_factor
2638 END IF
2639 IF (use_log_fallback) THEN
2640 CALL periodic_image_log_weights(image_coords, distances, log_weights)
2641 max_log_weight = maxval(log_weights)
2642 cpassert(max_log_weight > -0.5_dp*huge(1.0_dp))
2643 normalized_weights = 0.0_dp
2644 WHERE (log_weights > -0.5_dp*huge(1.0_dp))
2645 normalized_weights = exp(log_weights - max_log_weight)
2646 END WHERE
2647 normalization = sum(normalized_weights)
2648 cpassert(normalization > 0.0_dp)
2649 normalized_weights = normalized_weights/normalization
2650 RETURN
2651 END IF
2652
2653 screen_weight = reference_weight*screen_factor
2654 raw_weights = 0.0_dp
2655 raw_weights(nearest_image) = reference_weight
2656 DO iimage = 1, image_count
2657 IF (iimage == nearest_image) cycle
2658 pair = image_coords(:, iimage) - image_coords(:, nearest_image)
2659 rab = sqrt(sum(pair**2))
2660 IF (rab <= layout_tol) cycle
2661 mu = max(-1.0_dp, min(1.0_dp, &
2662 (distances(iimage) - distances(nearest_image))/rab))
2663 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2664 IF (factor <= tiny(1.0_dp) .OR. factor <= screen_weight) cycle
2665 product = factor
2666
2667 DO jimage = 1, image_count
2668 IF (jimage == iimage .OR. jimage == nearest_image) cycle
2669 pair = image_coords(:, iimage) - image_coords(:, jimage)
2670 rab = sqrt(sum(pair**2))
2671 IF (rab <= layout_tol) cycle
2672 mu = max(-1.0_dp, min(1.0_dp, &
2673 (distances(iimage) - distances(jimage))/rab))
2674 factor = 0.5_dp*(1.0_dp - becke_shape(mu))
2675 IF (factor <= tiny(1.0_dp) .OR. product <= screen_weight/factor) THEN
2676 product = 0.0_dp
2677 EXIT
2678 END IF
2679 product = product*factor
2680 END DO
2681 raw_weights(iimage) = product
2682 END DO
2683
2684 normalization = sum(raw_weights)
2685 cpassert(normalization > 0.0_dp)
2686 normalized_weights = raw_weights/normalization
2687
2688 END SUBROUTINE periodic_image_normalized_weights
2689
2690! **************************************************************************************************
2691!> \brief Build the image coordinates shared by all points of one target-atom block.
2692!> \param atom_coords ...
2693!> \param cell ...
2694!> \param target_atom ...
2695!> \param image_periodicity optional image directions used by the auxiliary-grid partition
2696!> \param image_coords ...
2697!> \param target_image ...
2698! **************************************************************************************************
2699 SUBROUTINE build_periodic_atom_image_layout(atom_coords, cell, target_atom, image_periodicity, &
2700 image_coords, target_image)
2701 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2702 TYPE(cell_type), POINTER :: cell
2703 INTEGER, INTENT(IN) :: target_atom
2704 INTEGER, DIMENSION(3), INTENT(IN), OPTIONAL :: image_periodicity
2705 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
2706 INTENT(OUT) :: image_coords
2707 INTEGER, INTENT(OUT) :: target_image
2708
2709 INTEGER :: base_shift(3), iatom, idir, image_count, &
2710 image_index, ix, ix_max, ix_min, iy, &
2711 iy_max, iy_min, iz, iz_max, iz_min, &
2712 natom
2713 INTEGER, DIMENSION(3) :: partition_periodicity
2714 REAL(kind=dp), DIMENSION(3) :: fractional
2715
2716 natom = SIZE(atom_coords, 2)
2717 cpassert(ASSOCIATED(cell))
2718 cpassert(target_atom >= 1 .AND. target_atom <= natom)
2719 cpassert(SIZE(atom_coords, 1) == 3)
2720 partition_periodicity = cell%perd
2721 IF (PRESENT(image_periodicity)) partition_periodicity = image_periodicity
2722 cpassert(all(partition_periodicity == 0 .OR. partition_periodicity == 1))
2723
2724 image_count = natom
2725 DO idir = 1, 3
2726 IF (partition_periodicity(idir) == 1) THEN
2727 image_count = image_count*(2*periodic_partition_image_shell + 1)
2728 END IF
2729 END DO
2730 ALLOCATE (image_coords(3, image_count))
2731
2732 image_index = 0
2733 target_image = 0
2734 DO iatom = 1, natom
2735 fractional = matmul(cell%h_inv, &
2736 atom_coords(:, target_atom) - atom_coords(:, iatom))
2737 DO idir = 1, 3
2738 base_shift(idir) = partition_periodicity(idir)*nint(fractional(idir))
2739 END DO
2740 ix_min = base_shift(1) - periodic_partition_image_shell*partition_periodicity(1)
2741 ix_max = base_shift(1) + periodic_partition_image_shell*partition_periodicity(1)
2742 iy_min = base_shift(2) - periodic_partition_image_shell*partition_periodicity(2)
2743 iy_max = base_shift(2) + periodic_partition_image_shell*partition_periodicity(2)
2744 iz_min = base_shift(3) - periodic_partition_image_shell*partition_periodicity(3)
2745 iz_max = base_shift(3) + periodic_partition_image_shell*partition_periodicity(3)
2746 DO iz = iz_min, iz_max
2747 DO iy = iy_min, iy_max
2748 DO ix = ix_min, ix_max
2749 image_index = image_index + 1
2750 image_coords(:, image_index) = atom_coords(:, iatom) + &
2751 matmul(cell%hmat, real([ix, iy, iz], kind=dp))
2752 IF (iatom == target_atom .AND. all([ix, iy, iz] == 0)) THEN
2753 target_image = image_index
2754 END IF
2755 END DO
2756 END DO
2757 END DO
2758 END DO
2759 cpassert(image_index == image_count)
2760 cpassert(target_image > 0)
2761
2762 END SUBROUTINE build_periodic_atom_image_layout
2763
2764! **************************************************************************************************
2765!> \brief Return an image-complete periodic atom weight for a prebuilt image layout.
2766!> \param grid_point ...
2767!> \param image_coords ...
2768!> \param target_image ...
2769!> \param weight ...
2770! **************************************************************************************************
2771 SUBROUTINE periodic_atom_image_partition_from_layout(grid_point, image_coords, target_image, &
2772 weight)
2773 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
2774 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: image_coords
2775 INTEGER, INTENT(IN) :: target_image
2776 REAL(kind=dp), INTENT(OUT) :: weight
2777
2778 INTEGER :: iimage, image_count
2779 REAL(kind=dp), DIMENSION(SIZE(image_coords, 2)) :: distances, normalized_weights
2780
2781 image_count = SIZE(image_coords, 2)
2782 cpassert(SIZE(image_coords, 1) == 3)
2783 cpassert(target_image >= 1 .AND. target_image <= image_count)
2784 IF (image_count == 1) THEN
2785 weight = 1.0_dp
2786 RETURN
2787 END IF
2788
2789 weight = 0.0_dp
2790 DO iimage = 1, image_count
2791 distances(iimage) = sqrt(sum((grid_point - image_coords(:, iimage))**2))
2792 END DO
2793 CALL periodic_image_normalized_weights(image_coords, distances, normalized_weights)
2794 weight = normalized_weights(target_image)
2795
2796 END SUBROUTINE periodic_atom_image_partition_from_layout
2797
2798! **************************************************************************************************
2799!> \brief Return an image-complete periodic atom weight without derivatives.
2800!> \param grid_point ...
2801!> \param atom_coords ...
2802!> \param cell ...
2803!> \param target_atom ...
2804!> \param weight ...
2805!> \param image_periodicity optional image directions used by the auxiliary-grid partition
2806! **************************************************************************************************
2807 SUBROUTINE periodic_atom_image_partition_value(grid_point, atom_coords, cell, target_atom, &
2808 weight, image_periodicity)
2809 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
2810 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2811 TYPE(cell_type), POINTER :: cell
2812 INTEGER, INTENT(IN) :: target_atom
2813 REAL(kind=dp), INTENT(OUT) :: weight
2814 INTEGER, DIMENSION(3), INTENT(IN), OPTIONAL :: image_periodicity
2815
2816 INTEGER :: target_image
2817 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: image_coords
2818
2819 CALL build_periodic_atom_image_layout(atom_coords, cell, target_atom, image_periodicity, &
2820 image_coords, target_image)
2821 CALL periodic_atom_image_partition_from_layout(grid_point, image_coords, target_image, weight)
2822 DEALLOCATE (image_coords)
2823
2824 END SUBROUTINE periodic_atom_image_partition_value
2825
2826! **************************************************************************************************
2827!> \brief Return the smooth weight of one reference-cell atom in an image-complete periodic
2828!> atom partition. Each atom's fixed image shell is centered on the image nearest to the
2829!> target atom. The list is constant over the target atom block and covariant under lattice
2830!> translations of individual atoms.
2831!> \param grid_point ...
2832!> \param atom_coords ...
2833!> \param cell ...
2834!> \param target_atom ...
2835!> \param weight ...
2836!> \param dweight_datom optional derivative at fixed grid point with respect to base atoms
2837!> \param dweight_dstrain optional derivative due to lattice-image translations
2838!> \param image_periodicity optional image directions used by the auxiliary-grid partition
2839! **************************************************************************************************
2840 SUBROUTINE periodic_atom_image_partition(grid_point, atom_coords, cell, target_atom, weight, &
2841 dweight_datom, dweight_dstrain, image_periodicity)
2842 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
2843 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2844 TYPE(cell_type), POINTER :: cell
2845 INTEGER, INTENT(IN) :: target_atom
2846 REAL(kind=dp), INTENT(OUT) :: weight
2847 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT), &
2848 OPTIONAL :: dweight_datom, dweight_dstrain
2849 INTEGER, DIMENSION(3), INTENT(IN), OPTIONAL :: image_periodicity
2850
2851 INTEGER :: active_count, active_image, active_index, base_shift(3), candidate_image, iatom, &
2852 idir, iimage, image_count, image_index, ix, ix_max, ix_min, iy, iy_max, iy_min, iz, &
2853 iz_max, iz_min, jatom, jdir, jimage, natom, target_image
2854 INTEGER, ALLOCATABLE, DIMENSION(:) :: active_images, image_atom
2855 INTEGER, ALLOCATABLE, DIMENSION(:, :) :: image_shift
2856 INTEGER, DIMENSION(3) :: partition_periodicity
2857 LOGICAL :: need_derivatives
2858 REAL(kind=dp) :: dist_diff, ds_dmu, mu, mu_raw, &
2859 one_minus_switch, rab, switch
2860 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: distances, normalized_weights
2861 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: image_coords, mean_log_derivative, &
2862 rvecs, target_log_derivative, &
2863 unit_rvecs
2864 REAL(kind=dp), DIMENSION(3) :: dmu_image_i, dmu_image_j, ds_image_i, ds_image_j, fractional, &
2865 image_translation, image_weight_derivative, pair, unit_pair
2866
2867 natom = SIZE(atom_coords, 2)
2868 cpassert(ASSOCIATED(cell))
2869 cpassert(target_atom >= 1 .AND. target_atom <= natom)
2870 cpassert(SIZE(atom_coords, 1) == 3)
2871 need_derivatives = PRESENT(dweight_datom)
2872 cpassert(PRESENT(dweight_dstrain) .EQV. need_derivatives)
2873 IF (.NOT. need_derivatives) THEN
2874 CALL periodic_atom_image_partition_value( &
2875 grid_point, atom_coords, cell, target_atom, weight, image_periodicity)
2876 RETURN
2877 END IF
2878 IF (need_derivatives) THEN
2879 cpassert(SIZE(dweight_datom, 1) == 3)
2880 cpassert(SIZE(dweight_datom, 2) == natom)
2881 cpassert(SIZE(dweight_dstrain, 1) == 3)
2882 cpassert(SIZE(dweight_dstrain, 2) == 3)
2883 dweight_datom = 0.0_dp
2884 dweight_dstrain = 0.0_dp
2885 END IF
2886 weight = 0.0_dp
2887 partition_periodicity = cell%perd
2888 IF (PRESENT(image_periodicity)) partition_periodicity = image_periodicity
2889 cpassert(all(partition_periodicity == 0 .OR. partition_periodicity == 1))
2890
2891 image_count = 1
2892 DO idir = 1, 3
2893 IF (partition_periodicity(idir) == 1) THEN
2894 image_count = image_count*(2*periodic_partition_image_shell + 1)
2895 END IF
2896 END DO
2897 image_count = natom*image_count
2898 ALLOCATE (active_images(image_count), image_atom(image_count), image_shift(3, image_count), &
2899 image_coords(3, image_count), distances(image_count), &
2900 normalized_weights(image_count), &
2901 rvecs(3, image_count), unit_rvecs(3, image_count), &
2902 mean_log_derivative(3, image_count), &
2903 target_log_derivative(3, image_count))
2904
2905 image_index = 0
2906 target_image = 0
2907 DO iatom = 1, natom
2908 fractional = matmul(cell%h_inv, &
2909 atom_coords(:, target_atom) - atom_coords(:, iatom))
2910 DO idir = 1, 3
2911 base_shift(idir) = partition_periodicity(idir)*nint(fractional(idir))
2912 END DO
2913 ix_min = base_shift(1) - periodic_partition_image_shell*partition_periodicity(1)
2914 ix_max = base_shift(1) + periodic_partition_image_shell*partition_periodicity(1)
2915 iy_min = base_shift(2) - periodic_partition_image_shell*partition_periodicity(2)
2916 iy_max = base_shift(2) + periodic_partition_image_shell*partition_periodicity(2)
2917 iz_min = base_shift(3) - periodic_partition_image_shell*partition_periodicity(3)
2918 iz_max = base_shift(3) + periodic_partition_image_shell*partition_periodicity(3)
2919 DO iz = iz_min, iz_max
2920 DO iy = iy_min, iy_max
2921 DO ix = ix_min, ix_max
2922 image_index = image_index + 1
2923 image_atom(image_index) = iatom
2924 image_shift(:, image_index) = [ix, iy, iz]
2925 image_coords(:, image_index) = atom_coords(:, iatom) + &
2926 matmul(cell%hmat, real([ix, iy, iz], kind=dp))
2927 IF (iatom == target_atom .AND. all(image_shift(:, image_index) == 0)) THEN
2928 target_image = image_index
2929 END IF
2930 END DO
2931 END DO
2932 END DO
2933 END DO
2934 cpassert(image_index == image_count)
2935 DO iimage = 1, image_count
2936 rvecs(:, iimage) = grid_point - image_coords(:, iimage)
2937 distances(iimage) = sqrt(sum(rvecs(:, iimage)**2))
2938 IF (distances(iimage) > layout_tol) THEN
2939 unit_rvecs(:, iimage) = rvecs(:, iimage)/distances(iimage)
2940 ELSE
2941 unit_rvecs(:, iimage) = 0.0_dp
2942 END IF
2943 END DO
2944
2945 CALL periodic_image_normalized_weights(image_coords, distances, normalized_weights)
2946 weight = normalized_weights(target_image)
2947
2948 IF (need_derivatives .AND. weight > 0.0_dp) THEN
2949 active_count = 0
2950 DO iimage = 1, image_count
2951 IF (normalized_weights(iimage) > 0.0_dp) THEN
2952 active_count = active_count + 1
2953 active_images(active_count) = iimage
2954 END IF
2955 END DO
2956 cpassert(active_count > 0)
2957 ! Revisit the image pairs after normalization to contract the logarithmic
2958 ! weight derivatives directly, avoiding both an image_count-squared tensor
2959 ! and pairs for which both screened normalized weights vanish.
2960 mean_log_derivative = 0.0_dp
2961 target_log_derivative = 0.0_dp
2962 DO active_index = 1, active_count
2963 active_image = active_images(active_index)
2964 DO candidate_image = 1, image_count
2965 IF (candidate_image == active_image) cycle
2966 ! A pair of active images is visited by the lower-index image only.
2967 IF (normalized_weights(candidate_image) > 0.0_dp .AND. &
2968 candidate_image < active_image) cycle
2969 iimage = min(active_image, candidate_image)
2970 jimage = max(active_image, candidate_image)
2971 pair = image_coords(:, iimage) - image_coords(:, jimage)
2972 rab = sqrt(sum(pair**2))
2973 IF (rab <= layout_tol) cycle
2974 unit_pair = pair/rab
2975 dist_diff = distances(iimage) - distances(jimage)
2976 mu_raw = dist_diff/rab
2977 mu = max(-1.0_dp, min(1.0_dp, mu_raw))
2978 switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2979 one_minus_switch = 1.0_dp - switch
2980 IF (abs(mu_raw) < 1.0_dp) THEN
2981 ds_dmu = -0.5_dp*becke_shape_derivative(mu)
2982 ELSE
2983 ds_dmu = 0.0_dp
2984 END IF
2985 IF (abs(ds_dmu) <= 0.0_dp .OR. switch <= tiny(1.0_dp) .OR. &
2986 one_minus_switch <= tiny(1.0_dp)) cycle
2987 dmu_image_i = (-unit_rvecs(:, iimage)*rab - dist_diff*unit_pair)/rab**2
2988 dmu_image_j = (unit_rvecs(:, jimage)*rab + dist_diff*unit_pair)/rab**2
2989 ds_image_i = ds_dmu*dmu_image_i
2990 ds_image_j = ds_dmu*dmu_image_j
2991
2992 mean_log_derivative(:, iimage) = mean_log_derivative(:, iimage) + &
2993 normalized_weights(iimage)*ds_image_i/switch - &
2994 normalized_weights(jimage)*ds_image_i/one_minus_switch
2995 mean_log_derivative(:, jimage) = mean_log_derivative(:, jimage) + &
2996 normalized_weights(iimage)*ds_image_j/switch - &
2997 normalized_weights(jimage)*ds_image_j/one_minus_switch
2998 IF (target_image == iimage) THEN
2999 target_log_derivative(:, iimage) = &
3000 target_log_derivative(:, iimage) + ds_image_i/switch
3001 target_log_derivative(:, jimage) = &
3002 target_log_derivative(:, jimage) + ds_image_j/switch
3003 ELSE IF (target_image == jimage) THEN
3004 target_log_derivative(:, iimage) = &
3005 target_log_derivative(:, iimage) - ds_image_i/one_minus_switch
3006 target_log_derivative(:, jimage) = &
3007 target_log_derivative(:, jimage) - ds_image_j/one_minus_switch
3008 END IF
3009 END DO
3010 END DO
3011 DO jimage = 1, image_count
3012 image_weight_derivative = weight*(target_log_derivative(:, jimage) - &
3013 mean_log_derivative(:, jimage))
3014 jatom = image_atom(jimage)
3015 dweight_datom(:, jatom) = dweight_datom(:, jatom) + &
3016 image_weight_derivative
3017 image_translation = matmul( &
3018 cell%hmat, real(image_shift(:, jimage), kind=dp))
3019 DO idir = 1, 3
3020 DO jdir = 1, 3
3021 dweight_dstrain(idir, jdir) = dweight_dstrain(idir, jdir) + &
3022 image_weight_derivative(idir)* &
3023 image_translation(jdir)
3024 END DO
3025 END DO
3026 END DO
3027 END IF
3028
3029 DEALLOCATE (active_images, image_atom, image_shift, image_coords, distances, &
3030 normalized_weights, rvecs, unit_rvecs, mean_log_derivative, &
3031 target_log_derivative)
3032
3033 END SUBROUTINE periodic_atom_image_partition
3034
3035! **************************************************************************************************
3036!> \brief Build smooth atom weights and their atom/cell deformation derivatives.
3037!> \param grid_point ...
3038!> \param atom_coords ...
3039!> \param cell ...
3040!> \param weights ...
3041!> \param included ...
3042!> \param dweights_datom ...
3043!> \param dweights_dstrain ...
3044! **************************************************************************************************
3045 SUBROUTINE skala_gpw_smooth_partition_derivatives(grid_point, atom_coords, cell, &
3046 weights, included, dweights_datom, &
3047 dweights_dstrain)
3048 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
3049 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
3050 TYPE(cell_type), POINTER :: cell
3051 REAL(kind=dp), DIMENSION(:), INTENT(OUT) :: weights
3052 LOGICAL, DIMENSION(:), INTENT(OUT) :: included
3053 REAL(kind=dp), DIMENSION(:, :, :), INTENT(OUT) :: dweights_datom, dweights_dstrain
3054
3055 INTEGER :: iatom, idir, jatom, jdir, natom
3056 REAL(kind=dp) :: dist_diff, ds_dmu, included_sum, mu, &
3057 mu_raw, one_minus_switch, rab, rsum, &
3058 switch, total
3059 REAL(kind=dp), DIMENSION(3) :: dmu_atom_i, dmu_atom_j, ds_atom_i, &
3060 ds_atom_j, pair, unit_pair
3061 REAL(kind=dp), DIMENSION(3, 3) :: dmu_strain, ds_strain, mean_strain
3062 REAL(kind=dp), DIMENSION(3, SIZE(atom_coords, 2), & SIZE(atom_coords, 2)) :: log_weight_atom
3063 REAL(kind=dp), DIMENSION(3, SIZE(atom_coords, 2)) :: mean_atom, partition_atom_coords, rvecs, &
3064 unit_rvecs
3065 REAL(kind=dp), &
3066 DIMENSION(3, 3, SIZE(atom_coords, 2)) :: log_weight_strain
3067 REAL(kind=dp), DIMENSION(SIZE(atom_coords, 2)) :: distances, normalized_weights, &
3068 raw_weights
3069
3070 natom = SIZE(atom_coords, 2)
3071 cpassert(SIZE(weights) == natom)
3072 cpassert(SIZE(included) == natom)
3073 cpassert(SIZE(dweights_datom, 1) == 3)
3074 cpassert(SIZE(dweights_datom, 2) == natom)
3075 cpassert(SIZE(dweights_datom, 3) == natom)
3076 cpassert(SIZE(dweights_dstrain, 1) == 3)
3077 cpassert(SIZE(dweights_dstrain, 2) == 3)
3078 cpassert(SIZE(dweights_dstrain, 3) == natom)
3079
3080 weights = 0.0_dp
3081 included = .false.
3082 dweights_datom = 0.0_dp
3083 dweights_dstrain = 0.0_dp
3084 raw_weights = 1.0_dp
3085 log_weight_atom = 0.0_dp
3086 log_weight_strain = 0.0_dp
3087
3088 DO iatom = 1, natom
3089 partition_atom_coords(:, iatom) = &
3090 nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
3091 rvecs(:, iatom) = grid_point - partition_atom_coords(:, iatom)
3092 distances(iatom) = sqrt(sum(rvecs(:, iatom)**2))
3093 IF (distances(iatom) > layout_tol) THEN
3094 unit_rvecs(:, iatom) = rvecs(:, iatom)/distances(iatom)
3095 ELSE
3096 unit_rvecs(:, iatom) = 0.0_dp
3097 END IF
3098 END DO
3099
3100 DO iatom = 1, natom - 1
3101 DO jatom = iatom + 1, natom
3102 pair = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
3103 rab = sqrt(sum(pair**2))
3104 IF (rab <= layout_tol) cycle
3105 unit_pair = pair/rab
3106 dist_diff = distances(iatom) - distances(jatom)
3107 mu_raw = dist_diff/rab
3108 mu = max(-1.0_dp, min(1.0_dp, mu_raw))
3109 switch = 0.5_dp*(1.0_dp - becke_shape(mu))
3110 one_minus_switch = 1.0_dp - switch
3111
3112 IF (abs(mu_raw) < 1.0_dp) THEN
3113 ds_dmu = -0.5_dp*becke_shape_derivative(mu)
3114 ELSE
3115 ds_dmu = 0.0_dp
3116 END IF
3117 IF (abs(ds_dmu) > 0.0_dp .AND. switch > tiny(1.0_dp) .AND. &
3118 one_minus_switch > tiny(1.0_dp)) THEN
3119 dmu_atom_i = (-unit_rvecs(:, iatom)*rab - dist_diff*unit_pair)/rab**2
3120 dmu_atom_j = (unit_rvecs(:, jatom)*rab + dist_diff*unit_pair)/rab**2
3121 ds_atom_i = ds_dmu*dmu_atom_i
3122 ds_atom_j = ds_dmu*dmu_atom_j
3123 log_weight_atom(:, iatom, iatom) = &
3124 log_weight_atom(:, iatom, iatom) + ds_atom_i/switch
3125 log_weight_atom(:, iatom, jatom) = &
3126 log_weight_atom(:, iatom, jatom) - ds_atom_i/one_minus_switch
3127 log_weight_atom(:, jatom, iatom) = &
3128 log_weight_atom(:, jatom, iatom) + ds_atom_j/switch
3129 log_weight_atom(:, jatom, jatom) = &
3130 log_weight_atom(:, jatom, jatom) - ds_atom_j/one_minus_switch
3131
3132 DO idir = 1, 3
3133 DO jdir = 1, 3
3134 dmu_strain(idir, jdir) = &
3135 ((unit_rvecs(idir, iatom)*rvecs(jdir, iatom) - &
3136 unit_rvecs(idir, jatom)*rvecs(jdir, jatom))*rab - &
3137 dist_diff*unit_pair(idir)*pair(jdir))/rab**2
3138 END DO
3139 END DO
3140 ds_strain = ds_dmu*dmu_strain
3141 log_weight_strain(:, :, iatom) = &
3142 log_weight_strain(:, :, iatom) + ds_strain/switch
3143 log_weight_strain(:, :, jatom) = &
3144 log_weight_strain(:, :, jatom) - ds_strain/one_minus_switch
3145 END IF
3146
3147 raw_weights(iatom) = raw_weights(iatom)*switch
3148 raw_weights(jatom) = raw_weights(jatom)*one_minus_switch
3149 END DO
3150 END DO
3151
3152 total = sum(raw_weights)
3153 IF (total > 0.0_dp) THEN
3154 normalized_weights = raw_weights/total
3155 included = normalized_weights > smooth_partition_eps
3156 ELSE
3157 rsum = huge(1.0_dp)
3158 jatom = 1
3159 DO iatom = 1, natom
3160 IF (distances(iatom) < rsum) THEN
3161 rsum = distances(iatom)
3162 jatom = iatom
3163 END IF
3164 END DO
3165 included(jatom) = .true.
3166 weights(jatom) = 1.0_dp
3167 RETURN
3168 END IF
3169
3170 included_sum = sum(raw_weights, mask=included)
3171 IF (included_sum <= 0.0_dp) THEN
3172 rsum = huge(1.0_dp)
3173 jatom = 1
3174 DO iatom = 1, natom
3175 IF (distances(iatom) < rsum) THEN
3176 rsum = distances(iatom)
3177 jatom = iatom
3178 END IF
3179 END DO
3180 included = .false.
3181 included(jatom) = .true.
3182 weights = 0.0_dp
3183 weights(jatom) = 1.0_dp
3184 RETURN
3185 END IF
3186
3187 DO iatom = 1, natom
3188 IF (included(iatom)) weights(iatom) = raw_weights(iatom)/included_sum
3189 END DO
3190
3191 mean_atom = 0.0_dp
3192 mean_strain = 0.0_dp
3193 DO iatom = 1, natom
3194 IF (.NOT. included(iatom)) cycle
3195 mean_strain = mean_strain + weights(iatom)*log_weight_strain(:, :, iatom)
3196 DO jatom = 1, natom
3197 mean_atom(:, jatom) = mean_atom(:, jatom) + &
3198 weights(iatom)*log_weight_atom(:, jatom, iatom)
3199 END DO
3200 END DO
3201
3202 DO iatom = 1, natom
3203 IF (.NOT. included(iatom)) cycle
3204 dweights_dstrain(:, :, iatom) = &
3205 weights(iatom)*(log_weight_strain(:, :, iatom) - mean_strain)
3206 DO jatom = 1, natom
3207 dweights_datom(:, jatom, iatom) = &
3208 weights(iatom)*(log_weight_atom(:, jatom, iatom) - mean_atom(:, jatom))
3209 END DO
3210 END DO
3211
3212 END SUBROUTINE skala_gpw_smooth_partition_derivatives
3213
3214! **************************************************************************************************
3215!> \brief Smoothly suppress a sparse atom row's internal quadrature weight at the layout cutoff.
3216!> \param weight normalized smooth-partition weight
3217!> \return scale in the interval [0, 1]
3218! **************************************************************************************************
3219 PURE FUNCTION smooth_partition_atomic_weight_scale(weight) RESULT(scale)
3220 REAL(kind=dp), INTENT(IN) :: weight
3221 REAL(kind=dp) :: scale
3222
3223 REAL(kind=dp) :: x
3224
3225 IF (weight <= smooth_partition_eps) THEN
3226 scale = 0.0_dp
3227 ELSE IF (weight >= smooth_partition_taper_end) THEN
3228 scale = 1.0_dp
3229 ELSE
3230 x = (weight - smooth_partition_eps)/ &
3231 (smooth_partition_taper_end - smooth_partition_eps)
3232 scale = x**3*(10.0_dp + x*(-15.0_dp + 6.0_dp*x))
3233 END IF
3234
3235 END FUNCTION smooth_partition_atomic_weight_scale
3236
3237! **************************************************************************************************
3238!> \brief Derivative of the sparse atom-row quadrature taper with respect to partition weight.
3239!> \param weight normalized smooth-partition weight
3240!> \return derivative of the internal quadrature-weight scale
3241! **************************************************************************************************
3242 PURE FUNCTION smooth_partition_atomic_weight_scale_derivative(weight) RESULT(derivative)
3243 REAL(kind=dp), INTENT(IN) :: weight
3244 REAL(kind=dp) :: derivative
3245
3246 REAL(kind=dp) :: x
3247
3248 IF (weight <= smooth_partition_eps .OR. weight >= smooth_partition_taper_end) THEN
3249 derivative = 0.0_dp
3250 ELSE
3251 x = (weight - smooth_partition_eps)/ &
3252 (smooth_partition_taper_end - smooth_partition_eps)
3253 derivative = 30.0_dp*x**2*(1.0_dp - x)**2/ &
3254 (smooth_partition_taper_end - smooth_partition_eps)
3255 END IF
3256
3257 END FUNCTION smooth_partition_atomic_weight_scale_derivative
3258
3259! **************************************************************************************************
3260!> \brief Becke fuzzy-cell shape function.
3261!> \param mu ...
3262!> \return ...
3263! **************************************************************************************************
3264 PURE FUNCTION becke_shape(mu) RESULT(val)
3265 REAL(kind=dp), INTENT(IN) :: mu
3266 REAL(kind=dp) :: val
3267
3268 INTEGER :: iter
3269
3270 val = mu
3271 DO iter = 1, 3
3272 val = 0.5_dp*val*(3.0_dp - val*val)
3273 END DO
3274
3275 END FUNCTION becke_shape
3276
3277! **************************************************************************************************
3278!> \brief Derivative of the Becke fuzzy-cell shape function.
3279!> \param mu ...
3280!> \return ...
3281! **************************************************************************************************
3282 PURE FUNCTION becke_shape_derivative(mu) RESULT(val)
3283 REAL(kind=dp), INTENT(IN) :: mu
3284 REAL(kind=dp) :: val
3285
3286 INTEGER :: iter
3287 REAL(kind=dp) :: x
3288
3289 x = mu
3290 val = 1.0_dp
3291 DO iter = 1, 3
3292 val = val*1.5_dp*(1.0_dp - x*x)
3293 x = 0.5_dp*x*(3.0_dp - x*x)
3294 END DO
3295
3296 END FUNCTION becke_shape_derivative
3297
3298! **************************************************************************************************
3299!> \brief Return the atom image nearest to a regular-grid point.
3300!> \param atom_coord ...
3301!> \param grid_point ...
3302!> \param cell ...
3303!> \return ...
3304! **************************************************************************************************
3305 FUNCTION nearest_atom_image_coordinate(atom_coord, grid_point, cell) RESULT(coord)
3306 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: atom_coord, grid_point
3307 TYPE(cell_type), POINTER :: cell
3308 REAL(kind=dp), DIMENSION(3) :: coord
3309
3310 REAL(kind=dp) :: dx, dy, dz
3311
3312 IF (cell%orthorhombic) THEN
3313 dx = atom_coord(1) - grid_point(1)
3314 dy = atom_coord(2) - grid_point(2)
3315 dz = atom_coord(3) - grid_point(3)
3316 dx = dx - cell%hmat(1, 1)*cell%perd(1)*anint(cell%h_inv(1, 1)*dx)
3317 dy = dy - cell%hmat(2, 2)*cell%perd(2)*anint(cell%h_inv(2, 2)*dy)
3318 dz = dz - cell%hmat(3, 3)*cell%perd(3)*anint(cell%h_inv(3, 3)*dz)
3319 coord = grid_point + [dx, dy, dz]
3320 ELSE
3321 coord = grid_point + pbc(grid_point, atom_coord, cell)
3322 END IF
3323
3324 END FUNCTION nearest_atom_image_coordinate
3325
3326! **************************************************************************************************
3327!> \brief Assign a grid point to the nearest periodic atom.
3328!> \param grid_point ...
3329!> \param atom_coords ...
3330!> \param cell ...
3331!> \return ...
3332! **************************************************************************************************
3333 FUNCTION nearest_atom(grid_point, atom_coords, cell) RESULT(owner)
3334 REAL(kind=dp), DIMENSION(3), INTENT(IN) :: grid_point
3335 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
3336 TYPE(cell_type), POINTER :: cell
3337 INTEGER :: owner
3338
3339 INTEGER :: iatom
3340 REAL(kind=dp) :: best_r2, dx, dy, dz, r2
3341 REAL(kind=dp), DIMENSION(3) :: rij
3342
3343 owner = 1
3344 best_r2 = huge(1.0_dp)
3345 IF (cell%orthorhombic) THEN
3346 DO iatom = 1, SIZE(atom_coords, 2)
3347 dx = grid_point(1) - atom_coords(1, iatom)
3348 dy = grid_point(2) - atom_coords(2, iatom)
3349 dz = grid_point(3) - atom_coords(3, iatom)
3350 dx = dx - cell%hmat(1, 1)*cell%perd(1)*anint(cell%h_inv(1, 1)*dx)
3351 dy = dy - cell%hmat(2, 2)*cell%perd(2)*anint(cell%h_inv(2, 2)*dy)
3352 dz = dz - cell%hmat(3, 3)*cell%perd(3)*anint(cell%h_inv(3, 3)*dz)
3353 r2 = dx*dx + dy*dy + dz*dz
3354 IF (r2 < best_r2) THEN
3355 best_r2 = r2
3356 owner = iatom
3357 END IF
3358 END DO
3359 ELSE
3360 DO iatom = 1, SIZE(atom_coords, 2)
3361 rij = pbc(grid_point, atom_coords(:, iatom), cell)
3362 r2 = sum(rij**2)
3363 IF (r2 < best_r2) THEN
3364 best_r2 = r2
3365 owner = iatom
3366 END IF
3367 END DO
3368 END IF
3369
3370 END FUNCTION nearest_atom
3371
3372END MODULE skala_gpw_features
3373
cell_types::pbc
Definition cell_types.F:105
torch_api::torch_tensor_from_array
Definition torch_api.F:44
torch_api::torch_tensor_reset_from_array
Definition torch_api.F:59
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
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
message_passing
Interface to the message passing library MPI.
Definition message_passing.F:27
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_types
Definition pw_types.F:24
skala_gpw_features
Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
Definition skala_gpw_features.F:11
skala_gpw_features::periodic_atom_image_partition_from_layout
subroutine, public periodic_atom_image_partition_from_layout(grid_point, image_coords, target_image, weight)
Return an image-complete periodic atom weight for a prebuilt image layout.
Definition skala_gpw_features.F:2773
skala_gpw_features::skala_gpw_atom_subchunk_layout
subroutine, public skala_gpw_atom_subchunk_layout(max_rows, atom_begin, atom_count, row_begin, row_count)
Build atom and row bounds for all atom-contiguous subchunks in one pass.
Definition skala_gpw_features.F:1871
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:1967
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::build_periodic_atom_image_layout
subroutine, public build_periodic_atom_image_layout(atom_coords, cell, target_atom, image_periodicity, image_coords, target_image)
Build the image coordinates shared by all points of one target-atom block.
Definition skala_gpw_features.F:2701
skala_gpw_features::cached_layout
type(skala_gpw_layout_cache_type), target, save cached_layout
Definition skala_gpw_features.F:181
skala_gpw_features::smooth_partition_atomic_weight_scale_derivative
pure real(kind=dp) function, public smooth_partition_atomic_weight_scale_derivative(weight)
Derivative of the sparse atom-row quadrature taper with respect to partition weight.
Definition skala_gpw_features.F:3244
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:1751
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:1834
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:204
skala_gpw_features::smooth_atom_partition
subroutine, public smooth_atom_partition(grid_point, atom_coords, cell, weights, partition_atom_coords, distances, pair_distances)
Build Becke-like smooth atom weights for one native-grid point.
Definition skala_gpw_features.F:2436
skala_gpw_features::skala_gpw_feature_build_atom_subchunk_bounds
subroutine, public skala_gpw_feature_build_atom_subchunk_bounds(parent, features, atom_begin, atom_count, row_begin, row_count, requires_grad)
Build an atom-contiguous subchunk feature bundle from precomputed bounds.
Definition skala_gpw_features.F:1996
skala_gpw_features::periodic_atom_image_partition
subroutine, public periodic_atom_image_partition(grid_point, atom_coords, cell, target_atom, weight, dweight_datom, dweight_dstrain, image_periodicity)
Return the smooth weight of one reference-cell atom in an image-complete periodic atom partition....
Definition skala_gpw_features.F:2842
skala_gpw_features::smooth_partition_atomic_weight_scale
pure real(kind=dp) function, public smooth_partition_atomic_weight_scale(weight)
Smoothly suppress a sparse atom row's internal quadrature weight at the layout cutoff.
Definition skala_gpw_features.F:3221
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:3048
skala_gpw_features::skala_gpw_atom_partition_smooth
integer, parameter, public skala_gpw_atom_partition_smooth
Definition skala_gpw_features.F:38
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:1843
torch_api::torch_tensor_narrow
subroutine, public torch_tensor_narrow(tensor, dim, start_index, length, result)
Creates a view of a contiguous tensor slice.
Definition torch_api.F:1431
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:1527
torch_api::torch_dict_create
subroutine, public torch_dict_create(dict)
Creates an empty Torch dictionary.
Definition torch_api.F:1726
torch_api::torch_tensor_expand_dim
subroutine, public torch_tensor_expand_dim(tensor, dim, extent, result)
Creates an expanded tensor view along one singleton dimension.
Definition torch_api.F:1394
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:1777
torch_api::torch_dict_clone
subroutine, public torch_dict_clone(source, target)
Clones a Torch dictionary.
Definition torch_api.F:1749
torch_api::torch_tensor_release
subroutine, public torch_tensor_release(tensor)
Releases a Torch tensor and all its ressources.
Definition torch_api.F:1702
xc_rho_set_types
contains the structure
Definition xc_rho_set_types.F:14
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
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
message_passing::mp_comm_type
Definition message_passing.F:179
particle_types::particle_type
Definition particle_types.F:35
pw_grid_types::pw_grid_type
Definition pw_grid_types.F:53
pw_types::pw_r3d_rs_type
Definition pw_types.F:62
skala_gpw_features::skala_gpw_feature_type
Definition skala_gpw_features.F:128
torch_api::torch_dict_type
Definition torch_api.F:34
torch_api::torch_tensor_type
Definition torch_api.F:29
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