Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Build SKALA TorchScript feature dictionaries from CP2K GPW real-space grids.
10 : ! **************************************************************************************************
11 : MODULE 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 :
42 : PUBLIC :: skala_gpw_atom_subchunk_count, skala_gpw_feature_build, &
43 : skala_gpw_feature_build_atom_subchunk, skala_gpw_feature_release, &
44 : skala_gpw_feature_type, skala_gpw_smooth_partition_derivatives
45 :
46 : TYPE skala_gpw_layout_cache_type
47 : INTEGER :: chunk_atom_begin = 1, chunk_atom_end = 0, &
48 : chunk_feature_begin = 1, &
49 : chunk_feature_count = 0, chunk_natom = 0, &
50 : natom = 0, nflat = 0, nflat_local = 0, &
51 : npoint = 0, nproc = 0, &
52 : atom_partition = skala_gpw_atom_partition_hard
53 : INTEGER, DIMENSION(2, 3) :: bo = 0, bounds = 0
54 : INTEGER, DIMENSION(3) :: npts = 0
55 : INTEGER, ALLOCATABLE, DIMENSION(:) :: dynamic_counts, dynamic_displs, &
56 : chunk_feature_counts, chunk_feature_displs, &
57 : chunk_grad_counts, chunk_grad_displs, &
58 : feature_counts, feature_displs, &
59 : feature_source_points, global_to_feature, &
60 : local_feature_counts, local_feature_offsets, &
61 : local_feature_points, local_feature_rows, &
62 : route_grad_return_recv_counts, &
63 : route_grad_return_recv_displs, &
64 : route_grad_return_send_counts, &
65 : route_grad_return_send_displs, &
66 : route_local_dest, chunk_return_positions, &
67 : route_point_recv_counts, &
68 : route_point_recv_displs, &
69 : route_point_send_counts, &
70 : route_point_send_displs, &
71 : route_send_local_rows
72 : INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
73 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes, chunk_atomic_grid_sizes, &
74 : chunk_feature_indices
75 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: local_feature_indices
76 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape, &
77 : chunk_atomic_grid_size_bound_shape
78 : TYPE(torch_dict_type) :: chunk_inputs
79 : TYPE(torch_dict_type) :: chunk_static_inputs
80 : TYPE(torch_dict_type) :: inputs
81 : TYPE(torch_dict_type) :: static_inputs
82 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
83 : TYPE(torch_tensor_type) :: atomic_grid_sizes_t
84 : TYPE(torch_tensor_type) :: atomic_grid_weights_t
85 : TYPE(torch_tensor_type) :: chunk_atomic_grid_size_bound_shape_t
86 : TYPE(torch_tensor_type) :: chunk_atomic_grid_sizes_t
87 : TYPE(torch_tensor_type) :: chunk_atomic_grid_weights_t
88 : TYPE(torch_tensor_type) :: chunk_coarse_0_atomic_coords_t
89 : TYPE(torch_tensor_type) :: chunk_density_t
90 : TYPE(torch_tensor_type) :: chunk_density_input_t
91 : TYPE(torch_tensor_type) :: chunk_feature_indices_t
92 : TYPE(torch_tensor_type) :: chunk_grad_t
93 : TYPE(torch_tensor_type) :: chunk_grad_input_t
94 : TYPE(torch_tensor_type) :: chunk_grid_coords_t
95 : TYPE(torch_tensor_type) :: chunk_grid_weights_t
96 : TYPE(torch_tensor_type) :: chunk_kin_t
97 : TYPE(torch_tensor_type) :: chunk_kin_input_t
98 : TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
99 : TYPE(torch_tensor_type) :: density_t
100 : TYPE(torch_tensor_type) :: grid_coords_t
101 : TYPE(torch_tensor_type) :: grid_weights_t
102 : TYPE(torch_tensor_type) :: grad_t
103 : TYPE(torch_tensor_type) :: kin_t
104 : TYPE(torch_tensor_type) :: local_feature_indices_t
105 : REAL(KIND=dp) :: dvol = 0.0_dp, weight_sum = 0.0_dp, &
106 : weight_sumsq = 0.0_dp
107 : REAL(KIND=dp), DIMENSION(3, 3) :: cell_hmat = 0.0_dp, dh = 0.0_dp
108 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, chunk_atomic_grid_weights, &
109 : chunk_grid_weights, grid_weights
110 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords, chunk_coarse_0_atomic_coords, &
111 : chunk_grid_coords, coarse_0_atomic_coords, &
112 : grid_coords
113 : LOGICAL :: active = .FALSE., has_weights = .FALSE., &
114 : chunk_dynamic_input_views_active = .FALSE., &
115 : chunk_dynamic_tensors_active = .FALSE., &
116 : chunk_inputs_active = .FALSE., &
117 : chunk_inputs_use_collapsed_rks = .FALSE., &
118 : chunk_static_tensors_active = .FALSE., &
119 : dynamic_tensors_active = .FALSE., &
120 : inputs_active = .FALSE., &
121 : static_tensors_active = .FALSE.
122 : END TYPE skala_gpw_layout_cache_type
123 :
124 : TYPE skala_gpw_feature_type
125 : INTEGER :: chunk_feature_count = 0, nflat = 0, &
126 : nflat_local = 0, &
127 : atom_partition = skala_gpw_atom_partition_hard
128 : TYPE(torch_dict_type) :: inputs
129 : TYPE(torch_tensor_type) :: atomic_grid_size_bound_shape_t
130 : TYPE(torch_tensor_type) :: atomic_grid_sizes_t
131 : TYPE(torch_tensor_type) :: atomic_grid_weights_t
132 : TYPE(torch_tensor_type) :: coarse_0_atomic_coords_t
133 : TYPE(torch_tensor_type) :: density_input_t
134 : TYPE(torch_tensor_type) :: density_t
135 : TYPE(torch_tensor_type) :: grad_t
136 : TYPE(torch_tensor_type) :: grad_input_t
137 : TYPE(torch_tensor_type) :: grid_coords_t
138 : TYPE(torch_tensor_type) :: grid_weights_t
139 : TYPE(torch_tensor_type) :: kin_input_t
140 : TYPE(torch_tensor_type) :: kin_t
141 : TYPE(torch_tensor_type) :: local_feature_indices_t
142 : INTEGER, ALLOCATABLE, DIMENSION(:) :: chunk_grad_counts, chunk_grad_displs, &
143 : local_feature_counts, local_feature_offsets, &
144 : local_feature_rows, &
145 : chunk_return_positions, &
146 : route_grad_return_recv_counts, &
147 : route_grad_return_recv_displs, &
148 : route_grad_return_send_counts, &
149 : route_grad_return_send_displs, &
150 : route_point_recv_counts, &
151 : route_point_recv_displs, &
152 : route_point_send_counts, &
153 : route_point_send_displs, &
154 : route_send_local_rows
155 : INTEGER, ALLOCATABLE, DIMENSION(:, :, :) :: feature_index
156 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:) :: atomic_grid_sizes
157 : INTEGER(KIND=int_8), ALLOCATABLE, DIMENSION(:, :) :: atomic_grid_size_bound_shape
158 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: atomic_grid_weights, grid_weights
159 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: chunk_density, chunk_kin, &
160 : coarse_0_atomic_coords, density, &
161 : grid_coords, kin
162 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: chunk_grad, grad
163 : REAL(KIND=dp) :: electron_count = 0.0_dp, &
164 : grid_weight_sum = 0.0_dp, &
165 : spin_moment = 0.0_dp
166 : LOGICAL :: active = .FALSE., owns_coordinate_tensor = .FALSE., &
167 : owns_grid_coordinate_tensor = .FALSE., &
168 : owns_weight_tensors = .FALSE., &
169 : owns_dynamic_tensors = .TRUE., &
170 : owns_inputs = .TRUE., &
171 : owns_static_tensors = .TRUE., &
172 : uses_atom_chunk_routing = .FALSE., &
173 : uses_atom_chunks = .FALSE., &
174 : uses_collapsed_rks_dynamic = .FALSE.
175 : END TYPE skala_gpw_feature_type
176 :
177 : TYPE(skala_gpw_layout_cache_type), SAVE :: cached_layout
178 :
179 : CONTAINS
180 :
181 : ! **************************************************************************************************
182 : !> \brief Build a flat SKALA molecular feature dictionary from a local GPW grid.
183 : !> \param features ...
184 : !> \param rho_set ...
185 : !> \param rho_r ...
186 : !> \param particle_set ...
187 : !> \param cell ...
188 : !> \param requires_grad ...
189 : !> \param weights ...
190 : !> \param requires_coordinate_grad ...
191 : !> \param requires_stress_grad ...
192 : !> \param use_atom_chunks ...
193 : !> \param route_atom_chunks ...
194 : !> \param atom_partition ...
195 : ! **************************************************************************************************
196 290 : SUBROUTINE skala_gpw_feature_build(features, rho_set, rho_r, particle_set, cell, &
197 : requires_grad, weights, requires_coordinate_grad, &
198 : requires_stress_grad, use_atom_chunks, route_atom_chunks, &
199 : atom_partition)
200 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
201 : TYPE(xc_rho_set_type), INTENT(IN) :: rho_set
202 : TYPE(pw_r3d_rs_type), DIMENSION(:), INTENT(IN) :: rho_r
203 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
204 : TYPE(cell_type), POINTER :: cell
205 : LOGICAL, INTENT(IN), OPTIONAL :: requires_grad
206 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
207 : LOGICAL, INTENT(IN), OPTIONAL :: requires_coordinate_grad, &
208 : requires_stress_grad, use_atom_chunks, &
209 : route_atom_chunks
210 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
211 :
212 : INTEGER :: handle, i, ipt, ispin, j, k, local_row, my_atom_partition, &
213 : ndynamic_local_per_point, nflat, nflat_local, nspins, phase_handle, real_base, row
214 : INTEGER, DIMENSION(2, 3) :: bo
215 : LOGICAL :: collapse_spin_dynamics, my_requires_coordinate_grad, my_requires_grad, &
216 : my_requires_stress_grad, my_route_atom_chunks, my_use_atom_chunks, &
217 : use_atom_chunk_protocol, use_atom_chunk_routing
218 290 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: global_dynamic, local_dynamic
219 290 : REAL(KIND=dp), DIMENSION(:, :, :), POINTER :: rho, rhoa, rhob, tau_a, tau_b, tau_total
220 3480 : TYPE(cp_3d_r_cp_type), DIMENSION(3) :: drho, drhoa, drhob
221 : TYPE(pw_grid_type), POINTER :: pw_grid
222 :
223 290 : CALL timeset("skala_gpw_feature_build", handle)
224 :
225 290 : my_requires_grad = .FALSE.
226 290 : IF (PRESENT(requires_grad)) my_requires_grad = requires_grad
227 290 : my_requires_coordinate_grad = .FALSE.
228 290 : IF (PRESENT(requires_coordinate_grad)) THEN
229 290 : my_requires_coordinate_grad = requires_coordinate_grad
230 : END IF
231 290 : my_requires_stress_grad = .FALSE.
232 290 : IF (PRESENT(requires_stress_grad)) my_requires_stress_grad = requires_stress_grad
233 290 : my_use_atom_chunks = .FALSE.
234 290 : IF (PRESENT(use_atom_chunks)) my_use_atom_chunks = use_atom_chunks
235 290 : my_route_atom_chunks = .FALSE.
236 290 : IF (PRESENT(route_atom_chunks)) my_route_atom_chunks = route_atom_chunks
237 290 : my_atom_partition = skala_gpw_atom_partition_hard
238 290 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
239 290 : IF (my_atom_partition /= skala_gpw_atom_partition_hard .AND. &
240 : my_atom_partition /= skala_gpw_atom_partition_smooth) THEN
241 0 : CALL cp_abort(__LOCATION__, "Unknown native SKALA atom-partition mode.")
242 : END IF
243 290 : CPASSERT(ASSOCIATED(cell))
244 290 : CPASSERT(ASSOCIATED(particle_set))
245 290 : CPASSERT(SIZE(rho_r) == 1 .OR. SIZE(rho_r) == 2)
246 290 : CPASSERT(ASSOCIATED(rho_r(1)%pw_grid))
247 290 : pw_grid => rho_r(1)%pw_grid
248 :
249 290 : nspins = SIZE(rho_r)
250 2900 : bo = pw_grid%bounds_local
251 290 : nflat_local = pw_grid%ngpts_local
252 :
253 290 : CALL timeset("skala_gpw_pre_release", phase_handle)
254 290 : CALL skala_gpw_feature_release(features)
255 290 : CALL timestop(phase_handle)
256 :
257 290 : CALL timeset("skala_gpw_layout_cache", phase_handle)
258 290 : CALL ensure_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition)
259 290 : CALL timestop(phase_handle)
260 290 : nflat = cached_layout%nflat
261 : use_atom_chunk_protocol = my_use_atom_chunks .AND. &
262 290 : .NOT. (my_requires_coordinate_grad .OR. my_requires_stress_grad)
263 290 : use_atom_chunk_routing = use_atom_chunk_protocol .AND. my_route_atom_chunks
264 290 : collapse_spin_dynamics = nspins == 1 .AND. use_atom_chunk_routing
265 290 : ndynamic_local_per_point = ndynamic_per_point
266 290 : IF (collapse_spin_dynamics) ndynamic_local_per_point = nrks_dynamic_per_point
267 870 : ALLOCATE (local_dynamic(ndynamic_local_per_point*nflat_local))
268 290 : local_dynamic = 0.0_dp
269 :
270 290 : CALL timeset("skala_gpw_pack_local", phase_handle)
271 290 : IF (nspins == 1) THEN
272 242 : CALL xc_rho_set_get(rho_set, rho=rho, drho=drho, tau=tau_total)
273 : ELSE
274 : CALL xc_rho_set_get(rho_set, rhoa=rhoa, rhob=rhob, drhoa=drhoa, drhob=drhob, &
275 48 : tau_a=tau_a, tau_b=tau_b)
276 : END IF
277 :
278 290 : local_row = 0
279 6076 : DO k = bo(1, 3), bo(2, 3)
280 139878 : DO j = bo(1, 2), bo(2, 2)
281 1983041 : DO i = bo(1, 1), bo(2, 1)
282 1843453 : local_row = local_row + 1
283 1843453 : real_base = ndynamic_local_per_point*(local_row - 1)
284 :
285 1977255 : IF (nspins == 1) THEN
286 1499203 : IF (collapse_spin_dynamics) THEN
287 105472 : local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
288 105472 : local_dynamic(real_base + 2) = 0.5_dp*drho(1)%array(i, j, k)
289 105472 : local_dynamic(real_base + 3) = 0.5_dp*drho(2)%array(i, j, k)
290 105472 : local_dynamic(real_base + 4) = 0.5_dp*drho(3)%array(i, j, k)
291 105472 : local_dynamic(real_base + 5) = 0.5_dp*tau_total(i, j, k)
292 : ELSE
293 1393731 : local_dynamic(real_base + 1) = 0.5_dp*rho(i, j, k)
294 1393731 : local_dynamic(real_base + 2) = 0.5_dp*rho(i, j, k)
295 4181193 : DO ispin = 1, 2
296 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 1) = &
297 2787462 : 0.5_dp*drho(1)%array(i, j, k)
298 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 2) = &
299 2787462 : 0.5_dp*drho(2)%array(i, j, k)
300 : local_dynamic(real_base + 2 + 3*(ispin - 1) + 3) = &
301 2787462 : 0.5_dp*drho(3)%array(i, j, k)
302 4181193 : local_dynamic(real_base + 8 + ispin) = 0.5_dp*tau_total(i, j, k)
303 : END DO
304 : END IF
305 : ELSE
306 344250 : local_dynamic(real_base + 1) = rhoa(i, j, k)
307 344250 : local_dynamic(real_base + 2) = rhob(i, j, k)
308 344250 : local_dynamic(real_base + 3) = drhoa(1)%array(i, j, k)
309 344250 : local_dynamic(real_base + 4) = drhoa(2)%array(i, j, k)
310 344250 : local_dynamic(real_base + 5) = drhoa(3)%array(i, j, k)
311 344250 : local_dynamic(real_base + 6) = drhob(1)%array(i, j, k)
312 344250 : local_dynamic(real_base + 7) = drhob(2)%array(i, j, k)
313 344250 : local_dynamic(real_base + 8) = drhob(3)%array(i, j, k)
314 344250 : local_dynamic(real_base + 9) = tau_a(i, j, k)
315 344250 : local_dynamic(real_base + 10) = tau_b(i, j, k)
316 : END IF
317 : END DO
318 : END DO
319 : END DO
320 290 : CALL timestop(phase_handle)
321 :
322 290 : CALL timeset("skala_gpw_copy_layout", phase_handle)
323 : CALL copy_cached_layout(features, my_requires_coordinate_grad .OR. my_requires_stress_grad, &
324 : my_requires_stress_grad .OR. &
325 : (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
326 520 : (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
327 290 : CALL timestop(phase_handle)
328 :
329 290 : IF (use_atom_chunk_routing) THEN
330 8 : CALL timeset("skala_gpw_route_dyn", phase_handle)
331 : CALL route_atom_chunk_dynamics(features, local_dynamic, pw_grid%para%group, &
332 8 : collapse_spin_dynamics)
333 8 : features%uses_atom_chunk_routing = .TRUE.
334 8 : features%uses_atom_chunks = .TRUE.
335 8 : CALL timestop(phase_handle)
336 : ELSE
337 846 : ALLOCATE (global_dynamic(ndynamic_per_point*cached_layout%npoint))
338 282 : CALL timeset("skala_gpw_allgatherv", phase_handle)
339 : CALL pw_grid%para%group%allgatherv(local_dynamic, global_dynamic, &
340 : cached_layout%dynamic_counts, &
341 282 : cached_layout%dynamic_displs)
342 282 : CALL timestop(phase_handle)
343 :
344 282 : CALL timeset("skala_gpw_reorder_dyn", phase_handle)
345 0 : ALLOCATE (features%density(nflat, 2), features%grad(nflat, 3, 2), &
346 1974 : features%kin(nflat, 2))
347 10194530 : features%density = 0.0_dp
348 30583590 : features%grad = 0.0_dp
349 10194530 : features%kin = 0.0_dp
350 :
351 5097124 : DO row = 1, nflat
352 5096842 : ipt = cached_layout%feature_source_points(row)
353 5096842 : real_base = ndynamic_per_point*(ipt - 1)
354 15290526 : features%density(row, :) = global_dynamic(real_base + 1:real_base + 2)
355 5096842 : features%grad(row, 1, 1) = global_dynamic(real_base + 3)
356 5096842 : features%grad(row, 2, 1) = global_dynamic(real_base + 4)
357 5096842 : features%grad(row, 3, 1) = global_dynamic(real_base + 5)
358 5096842 : features%grad(row, 1, 2) = global_dynamic(real_base + 6)
359 5096842 : features%grad(row, 2, 2) = global_dynamic(real_base + 7)
360 5096842 : features%grad(row, 3, 2) = global_dynamic(real_base + 8)
361 15290808 : features%kin(row, :) = global_dynamic(real_base + 9:real_base + 10)
362 : END DO
363 846 : CALL timestop(phase_handle)
364 : END IF
365 :
366 290 : CALL timeset("skala_gpw_feature_sums", phase_handle)
367 290 : IF (features%uses_atom_chunks) THEN
368 8 : features%electron_count = 0.0_dp
369 8 : features%spin_moment = 0.0_dp
370 8 : IF (features%chunk_feature_count > 0) THEN
371 8 : IF (features%uses_collapsed_rks_dynamic) THEN
372 : features%electron_count = SUM(2.0_dp*features%chunk_density(:, 1)* &
373 146742 : cached_layout%chunk_grid_weights)
374 : ELSE
375 : features%electron_count = SUM((features%chunk_density(:, 1) + &
376 : features%chunk_density(:, 2))* &
377 0 : cached_layout%chunk_grid_weights)
378 : features%spin_moment = SUM((features%chunk_density(:, 1) - &
379 : features%chunk_density(:, 2))* &
380 0 : cached_layout%chunk_grid_weights)
381 : END IF
382 : END IF
383 8 : CALL pw_grid%para%group%sum(features%electron_count)
384 8 : CALL pw_grid%para%group%sum(features%spin_moment)
385 : ELSE
386 : features%electron_count = SUM((features%density(:, 1) + features%density(:, 2))* &
387 5097124 : features%grid_weights)
388 : features%spin_moment = SUM((features%density(:, 1) - features%density(:, 2))* &
389 5097124 : features%grid_weights)
390 : END IF
391 5390600 : features%grid_weight_sum = SUM(features%grid_weights)
392 290 : CALL timestop(phase_handle)
393 :
394 290 : CALL timeset("skala_gpw_tensor_update", phase_handle)
395 290 : IF (use_atom_chunk_protocol .AND. .NOT. features%uses_atom_chunks) THEN
396 0 : IF (features%chunk_feature_count > 0) CALL extract_atom_chunk_dynamics(features)
397 0 : features%uses_atom_chunks = .TRUE.
398 : END IF
399 290 : IF (.NOT. features%uses_atom_chunks .OR. features%chunk_feature_count > 0) THEN
400 : CALL add_feature_tensors(features, my_requires_grad, my_requires_coordinate_grad, &
401 : my_requires_stress_grad, &
402 : features%uses_atom_chunks, &
403 : requires_weight_grad= &
404 : (my_atom_partition == skala_gpw_atom_partition_smooth .AND. &
405 520 : (my_requires_coordinate_grad .OR. my_requires_stress_grad)))
406 : ELSE
407 : ! This rank participates in atom-chunk communication but owns no model input rows.
408 0 : features%owns_coordinate_tensor = .FALSE.
409 0 : features%owns_grid_coordinate_tensor = .FALSE.
410 0 : features%owns_weight_tensors = .FALSE.
411 0 : features%owns_dynamic_tensors = .FALSE.
412 0 : features%owns_inputs = .FALSE.
413 0 : features%owns_static_tensors = .FALSE.
414 : END IF
415 290 : CALL timestop(phase_handle)
416 290 : features%active = .TRUE.
417 :
418 290 : IF (ALLOCATED(global_dynamic)) DEALLOCATE (global_dynamic)
419 290 : DEALLOCATE (local_dynamic)
420 290 : CALL timestop(handle)
421 :
422 2320 : END SUBROUTINE skala_gpw_feature_build
423 :
424 : ! **************************************************************************************************
425 : !> \brief Ensure that static grid-to-atom layout data is cached for the current grid/geometry.
426 : !> \param pw_grid ...
427 : !> \param particle_set ...
428 : !> \param cell ...
429 : !> \param weights ...
430 : !> \param atom_partition ...
431 : ! **************************************************************************************************
432 290 : SUBROUTINE ensure_layout_cache(pw_grid, particle_set, cell, weights, atom_partition)
433 : TYPE(pw_grid_type), POINTER :: pw_grid
434 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
435 : TYPE(cell_type), POINTER :: cell
436 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
437 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
438 :
439 : INTEGER :: my_atom_partition, phase_handle
440 : LOGICAL :: cache_matches
441 :
442 290 : my_atom_partition = skala_gpw_atom_partition_hard
443 290 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
444 290 : IF (PRESENT(weights)) THEN
445 290 : CALL timeset("skala_gpw_layout_match", phase_handle)
446 : cache_matches = layout_cache_matches(pw_grid, particle_set, cell, weights, &
447 290 : my_atom_partition)
448 290 : CALL timestop(phase_handle)
449 290 : IF (cache_matches) RETURN
450 128 : CALL timeset("skala_gpw_layout_rebuild", phase_handle)
451 128 : CALL rebuild_layout_cache(pw_grid, particle_set, cell, weights, my_atom_partition)
452 128 : CALL timestop(phase_handle)
453 : ELSE
454 0 : CALL timeset("skala_gpw_layout_match", phase_handle)
455 : cache_matches = layout_cache_matches(pw_grid, particle_set, cell, &
456 0 : atom_partition=my_atom_partition)
457 0 : CALL timestop(phase_handle)
458 0 : IF (cache_matches) RETURN
459 0 : CALL timeset("skala_gpw_layout_rebuild", phase_handle)
460 : CALL rebuild_layout_cache(pw_grid, particle_set, cell, &
461 0 : atom_partition=my_atom_partition)
462 0 : CALL timestop(phase_handle)
463 : END IF
464 :
465 : END SUBROUTINE ensure_layout_cache
466 :
467 : ! **************************************************************************************************
468 : !> \brief Check whether the current static layout cache can be reused.
469 : !> \param pw_grid ...
470 : !> \param particle_set ...
471 : !> \param cell ...
472 : !> \param weights ...
473 : !> \param atom_partition ...
474 : !> \return ...
475 : ! **************************************************************************************************
476 290 : FUNCTION layout_cache_matches(pw_grid, particle_set, cell, weights, atom_partition) RESULT(matches)
477 : TYPE(pw_grid_type), POINTER :: pw_grid
478 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
479 : TYPE(cell_type), POINTER :: cell
480 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
481 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
482 : LOGICAL :: matches
483 :
484 : INTEGER :: iatom, my_atom_partition
485 : LOGICAL :: weights_match
486 :
487 290 : my_atom_partition = skala_gpw_atom_partition_hard
488 290 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
489 290 : matches = .FALSE.
490 290 : IF (.NOT. cached_layout%active) RETURN
491 202 : IF (cached_layout%atom_partition /= my_atom_partition) RETURN
492 202 : IF (cached_layout%natom /= SIZE(particle_set)) RETURN
493 202 : IF (cached_layout%nflat_local /= pw_grid%ngpts_local) RETURN
494 202 : IF (cached_layout%nproc /= pw_grid%para%group%num_pe) RETURN
495 2020 : IF (ANY(cached_layout%bo /= pw_grid%bounds_local)) RETURN
496 2020 : IF (ANY(cached_layout%bounds /= pw_grid%bounds)) RETURN
497 808 : IF (ANY(cached_layout%npts /= pw_grid%npts)) RETURN
498 202 : IF (ABS(cached_layout%dvol - pw_grid%dvol) > layout_tol) RETURN
499 2262 : IF (ANY(ABS(cached_layout%dh - pw_grid%dh) > layout_tol)) RETURN
500 2262 : IF (ANY(ABS(cached_layout%cell_hmat - cell%hmat) > layout_tol)) RETURN
501 174 : IF (.NOT. ALLOCATED(cached_layout%atom_coords)) RETURN
502 :
503 510 : DO iatom = 1, SIZE(particle_set)
504 1542 : IF (ANY(ABS(cached_layout%atom_coords(:, iatom) - particle_set(iatom)%r) > layout_tol)) RETURN
505 : END DO
506 :
507 162 : IF (PRESENT(weights)) THEN
508 162 : weights_match = layout_weights_match(pw_grid, weights)
509 : ELSE
510 0 : weights_match = layout_weights_match(pw_grid)
511 : END IF
512 162 : IF (.NOT. weights_match) RETURN
513 :
514 290 : matches = .TRUE.
515 :
516 : END FUNCTION layout_cache_matches
517 :
518 : ! **************************************************************************************************
519 : !> \brief Check whether current optional integration weights match the cached static tensors.
520 : !> \param pw_grid ...
521 : !> \param weights ...
522 : !> \return ...
523 : ! **************************************************************************************************
524 162 : FUNCTION layout_weights_match(pw_grid, weights) RESULT(matches)
525 : TYPE(pw_grid_type), POINTER :: pw_grid
526 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
527 : LOGICAL :: matches
528 :
529 : LOGICAL :: has_weights
530 : REAL(KIND=dp) :: weight_sum, weight_sumsq
531 :
532 162 : matches = .FALSE.
533 : MARK_USED(pw_grid)
534 162 : IF (PRESENT(weights)) THEN
535 162 : CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
536 : ELSE
537 : CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
538 0 : weight_sumsq=weight_sumsq)
539 : END IF
540 :
541 162 : IF (cached_layout%has_weights .NEQV. has_weights) RETURN
542 162 : IF (ABS(cached_layout%weight_sum - weight_sum) > layout_tol) RETURN
543 162 : IF (ABS(cached_layout%weight_sumsq - weight_sumsq) > layout_tol) RETURN
544 :
545 162 : matches = .TRUE.
546 :
547 : END FUNCTION layout_weights_match
548 :
549 : ! **************************************************************************************************
550 : !> \brief Build the static SKALA layout cache.
551 : !> \param pw_grid ...
552 : !> \param particle_set ...
553 : !> \param cell ...
554 : !> \param weights ...
555 : !> \param atom_partition ...
556 : ! **************************************************************************************************
557 128 : SUBROUTINE rebuild_layout_cache(pw_grid, particle_set, cell, weights, atom_partition)
558 : TYPE(pw_grid_type), POINTER :: pw_grid
559 : TYPE(particle_type), DIMENSION(:), POINTER :: particle_set
560 : TYPE(cell_type), POINTER :: cell
561 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
562 : INTEGER, INTENT(IN), OPTIONAL :: atom_partition
563 :
564 : INTEGER :: feature_local, i, iatom, ipt, j, k, local_row, max_grid_size, max_local_features, &
565 : my_atom_partition, natom, nfeature_local, nflat, nflat_local, npoint, nproc, owner, pe, &
566 : pe_index, phase_handle, row, source_global, source_local, static_base
567 128 : INTEGER, ALLOCATABLE, DIMENSION(:) :: atom_offset, atom_position, chunk_atom_begin, &
568 128 : chunk_atom_end, cursor, feature_counts, feature_displs, global_owner, &
569 128 : global_source_points, local_feature_counts_tmp, local_owner, local_source_global, &
570 128 : local_source_points, point_counts, point_displs, static_counts, static_displs
571 : INTEGER, DIMENSION(2, 3) :: bo
572 : LOGICAL :: has_weights
573 : REAL(KIND=dp) :: base_weight, included_sum, &
574 : partition_weight, weight_sum, &
575 : weight_sumsq
576 128 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: distances, global_static, local_static, &
577 : partition_weights
578 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: atom_coords_pbc, atom_image_coords
579 : REAL(KIND=dp), DIMENSION(3) :: grid_point, owner_coord
580 :
581 128 : CALL release_layout_cache(cached_layout)
582 :
583 128 : my_atom_partition = skala_gpw_atom_partition_hard
584 128 : IF (PRESENT(atom_partition)) my_atom_partition = atom_partition
585 128 : natom = SIZE(particle_set)
586 1280 : bo = pw_grid%bounds_local
587 128 : nflat_local = pw_grid%ngpts_local
588 128 : nproc = pw_grid%para%group%num_pe
589 128 : pe_index = pw_grid%para%group%mepos + 1
590 :
591 128 : IF (PRESENT(weights)) THEN
592 128 : CALL weights_signature(weights, has_weights, weight_sum, weight_sumsq)
593 : ELSE
594 : CALL weights_signature(has_weights=has_weights, weight_sum=weight_sum, &
595 0 : weight_sumsq=weight_sumsq)
596 : END IF
597 :
598 128 : max_local_features = nflat_local
599 128 : IF (my_atom_partition == skala_gpw_atom_partition_smooth) THEN
600 102 : max_local_features = nflat_local*natom
601 : END IF
602 : ALLOCATE (local_owner(max_local_features), &
603 : local_source_points(max_local_features), &
604 : local_static(nstatic_per_point*max_local_features), &
605 : local_feature_counts_tmp(nflat_local), feature_counts(nproc), &
606 : feature_displs(nproc), point_counts(nproc), point_displs(nproc), &
607 : static_counts(nproc), static_displs(nproc), atom_coords_pbc(3, natom), &
608 2688 : atom_image_coords(3, natom), distances(natom), partition_weights(natom))
609 0 : ALLOCATE (cached_layout%feature_index(bo(1, 1):bo(2, 1), &
610 : bo(1, 2):bo(2, 2), &
611 640 : bo(1, 3):bo(2, 3)))
612 1451131 : cached_layout%feature_index = 0
613 128 : local_static = 0.0_dp
614 128 : local_feature_counts_tmp = 0
615 412 : DO iatom = 1, natom
616 412 : atom_coords_pbc(:, iatom) = pbc(particle_set(iatom)%r, cell, positive_range=.TRUE.)
617 : END DO
618 :
619 128 : CALL timeset("skala_gpw_layout_local", phase_handle)
620 128 : local_row = 0
621 128 : nfeature_local = 0
622 3118 : DO k = bo(1, 3), bo(2, 3)
623 86436 : DO j = bo(1, 2), bo(2, 2)
624 1451003 : DO i = bo(1, 1), bo(2, 1)
625 1364695 : local_row = local_row + 1
626 5458780 : grid_point = grid_coordinate(pw_grid, [i, j, k])
627 1364695 : base_weight = pw_grid%dvol
628 1364695 : IF (PRESENT(weights)) THEN
629 1364695 : IF (ASSOCIATED(weights)) base_weight = base_weight*weights%array(i, j, k)
630 : END IF
631 1364695 : cached_layout%feature_index(i, j, k) = local_row
632 :
633 1448013 : IF (my_atom_partition == skala_gpw_atom_partition_hard) THEN
634 987187 : owner = nearest_atom(grid_point, atom_coords_pbc, cell)
635 3948748 : owner_coord = atom_coords_pbc(:, owner)
636 987187 : nfeature_local = nfeature_local + 1
637 987187 : local_feature_counts_tmp(local_row) = 1
638 987187 : local_owner(nfeature_local) = owner
639 987187 : local_source_points(nfeature_local) = local_row
640 987187 : static_base = nstatic_per_point*(nfeature_local - 1)
641 3948748 : local_static(static_base + 1:static_base + 3) = grid_point
642 987187 : local_static(static_base + 4) = base_weight
643 987187 : local_static(static_base + 5) = base_weight
644 : ELSE
645 : CALL smooth_atom_partition(grid_point, atom_coords_pbc, cell, &
646 377508 : partition_weights, atom_image_coords, distances)
647 1132524 : included_sum = SUM(partition_weights, MASK=partition_weights > smooth_partition_eps)
648 377508 : IF (included_sum <= 0.0_dp) THEN
649 0 : owner = nearest_atom(grid_point, atom_coords_pbc, cell)
650 0 : partition_weights = 0.0_dp
651 0 : partition_weights(owner) = 1.0_dp
652 0 : included_sum = 1.0_dp
653 : END IF
654 1132524 : DO iatom = 1, natom
655 755016 : IF (partition_weights(iatom) <= smooth_partition_eps) CYCLE
656 753034 : partition_weight = partition_weights(iatom)/included_sum
657 753034 : nfeature_local = nfeature_local + 1
658 : local_feature_counts_tmp(local_row) = &
659 753034 : local_feature_counts_tmp(local_row) + 1
660 753034 : local_owner(nfeature_local) = iatom
661 753034 : local_source_points(nfeature_local) = local_row
662 753034 : static_base = nstatic_per_point*(nfeature_local - 1)
663 3012136 : local_static(static_base + 1:static_base + 3) = grid_point
664 753034 : local_static(static_base + 4) = base_weight*partition_weight
665 1132524 : local_static(static_base + 5) = base_weight
666 : END DO
667 : END IF
668 : END DO
669 : END DO
670 : END DO
671 128 : CALL timestop(phase_handle)
672 :
673 : ! SKALA groups all grid points by atom. This ordering is static while the
674 : ! grid, cell, atom positions, and optional integration weights are unchanged.
675 128 : CALL timeset("skala_gpw_layout_gather", phase_handle)
676 128 : CALL pw_grid%para%group%allgather(nflat_local, point_counts)
677 128 : CALL counts_to_displs(point_counts, point_displs)
678 384 : npoint = SUM(point_counts)
679 128 : CALL pw_grid%para%group%allgather(nfeature_local, feature_counts)
680 128 : CALL counts_to_displs(feature_counts, feature_displs)
681 384 : DO pe = 1, nproc
682 256 : static_counts(pe) = nstatic_per_point*feature_counts(pe)
683 384 : static_displs(pe) = nstatic_per_point*feature_displs(pe)
684 : END DO
685 384 : nflat = SUM(feature_counts)
686 : ALLOCATE (global_owner(nflat), global_source_points(nflat), &
687 1024 : global_static(nstatic_per_point*nflat), local_source_global(nfeature_local))
688 1740349 : DO feature_local = 1, nfeature_local
689 1740349 : local_source_global(feature_local) = point_displs(pe_index) + local_source_points(feature_local)
690 : END DO
691 : CALL pw_grid%para%group%allgatherv(local_owner(1:nfeature_local), global_owner, feature_counts, &
692 128 : feature_displs)
693 : CALL pw_grid%para%group%allgatherv(local_source_global, global_source_points, feature_counts, &
694 128 : feature_displs)
695 : CALL pw_grid%para%group%allgatherv(local_static(1:nstatic_per_point*nfeature_local), &
696 : global_static, static_counts, &
697 128 : static_displs)
698 128 : CALL timestop(phase_handle)
699 :
700 0 : ALLOCATE (cached_layout%chunk_feature_counts(nproc), &
701 0 : cached_layout%chunk_feature_displs(nproc), &
702 0 : cached_layout%chunk_grad_counts(nproc), cached_layout%chunk_grad_displs(nproc), &
703 0 : cached_layout%feature_counts(nproc), cached_layout%feature_displs(nproc), &
704 0 : cached_layout%dynamic_counts(nproc), cached_layout%dynamic_displs(nproc), &
705 0 : cached_layout%route_grad_return_recv_counts(nproc), &
706 0 : cached_layout%route_grad_return_recv_displs(nproc), &
707 0 : cached_layout%route_grad_return_send_counts(nproc), &
708 0 : cached_layout%route_grad_return_send_displs(nproc), &
709 0 : cached_layout%route_point_recv_counts(nproc), &
710 0 : cached_layout%route_point_recv_displs(nproc), &
711 0 : cached_layout%route_point_send_counts(nproc), &
712 0 : cached_layout%route_point_send_displs(nproc), &
713 0 : cached_layout%feature_source_points(nflat), &
714 0 : cached_layout%global_to_feature(npoint), cached_layout%atomic_grid_sizes(natom), &
715 0 : cached_layout%local_feature_counts(nflat_local), &
716 0 : cached_layout%local_feature_offsets(nflat_local + 1), &
717 0 : cached_layout%local_feature_rows(nfeature_local), &
718 0 : cached_layout%local_feature_points(nfeature_local), &
719 0 : cached_layout%local_feature_indices(nfeature_local), atom_offset(natom + 1), &
720 : atom_position(natom), chunk_atom_begin(nproc), chunk_atom_end(nproc), &
721 4480 : cursor(nflat_local))
722 384 : cached_layout%feature_counts(:) = feature_counts
723 384 : cached_layout%feature_displs(:) = feature_displs
724 384 : cached_layout%dynamic_counts(:) = ndynamic_per_point*point_counts
725 384 : cached_layout%dynamic_displs(:) = ndynamic_per_point*point_displs
726 412 : cached_layout%atomic_grid_sizes = 0_int_8
727 2729518 : cached_layout%global_to_feature = 0
728 1364823 : cached_layout%local_feature_counts(:) = local_feature_counts_tmp
729 128 : cached_layout%local_feature_offsets(1) = 1
730 1364823 : DO local_row = 1, nflat_local
731 : cached_layout%local_feature_offsets(local_row + 1) = &
732 : cached_layout%local_feature_offsets(local_row) + &
733 1364823 : cached_layout%local_feature_counts(local_row)
734 : END DO
735 1364823 : cursor(:) = cached_layout%local_feature_offsets(1:nflat_local)
736 :
737 128 : CALL timeset("skala_gpw_layout_atom_sort", phase_handle)
738 3480570 : DO ipt = 1, nflat
739 : cached_layout%atomic_grid_sizes(global_owner(ipt)) = &
740 3480570 : cached_layout%atomic_grid_sizes(global_owner(ipt)) + 1_int_8
741 : END DO
742 128 : atom_offset(1) = 1
743 412 : DO iatom = 1, natom
744 412 : atom_offset(iatom + 1) = atom_offset(iatom) + INT(cached_layout%atomic_grid_sizes(iatom))
745 : END DO
746 412 : DO iatom = 1, natom
747 412 : atom_position(iatom) = atom_offset(iatom)
748 : END DO
749 412 : max_grid_size = MAXVAL(INT(cached_layout%atomic_grid_sizes))
750 : CALL build_atom_chunks(cached_layout%atomic_grid_sizes, atom_offset, nproc, &
751 : chunk_atom_begin, chunk_atom_end, &
752 : cached_layout%chunk_feature_counts, &
753 128 : cached_layout%chunk_feature_displs)
754 384 : cached_layout%chunk_grad_counts(:) = ngrad_per_point*cached_layout%chunk_feature_counts
755 384 : cached_layout%chunk_grad_displs(:) = ngrad_per_point*cached_layout%chunk_feature_displs
756 128 : cached_layout%chunk_atom_begin = chunk_atom_begin(pe_index)
757 128 : cached_layout%chunk_atom_end = chunk_atom_end(pe_index)
758 128 : cached_layout%chunk_feature_begin = cached_layout%chunk_feature_displs(pe_index) + 1
759 128 : cached_layout%chunk_feature_count = cached_layout%chunk_feature_counts(pe_index)
760 : cached_layout%chunk_natom = cached_layout%chunk_atom_end - &
761 128 : cached_layout%chunk_atom_begin + 1
762 :
763 0 : ALLOCATE (cached_layout%grid_coords(3, nflat), cached_layout%grid_weights(nflat), &
764 0 : cached_layout%atomic_grid_weights(nflat), &
765 0 : cached_layout%coarse_0_atomic_coords(3, natom), &
766 0 : cached_layout%atomic_grid_size_bound_shape(0, max_grid_size), &
767 1152 : cached_layout%atom_coords(3, natom))
768 13921896 : cached_layout%grid_coords = 0.0_dp
769 3480570 : cached_layout%grid_weights = 0.0_dp
770 3480570 : cached_layout%atomic_grid_weights = 0.0_dp
771 1533694 : cached_layout%atomic_grid_size_bound_shape = 0_int_8
772 :
773 412 : DO iatom = 1, natom
774 1136 : cached_layout%atom_coords(:, iatom) = particle_set(iatom)%r
775 1264 : cached_layout%coarse_0_atomic_coords(:, iatom) = atom_coords_pbc(:, iatom)
776 : END DO
777 :
778 3480570 : DO ipt = 1, nflat
779 3480442 : owner = global_owner(ipt)
780 3480442 : row = atom_position(owner)
781 3480442 : atom_position(owner) = atom_position(owner) + 1
782 3480442 : source_global = global_source_points(ipt)
783 3480442 : cached_layout%feature_source_points(row) = source_global
784 3480442 : IF (cached_layout%global_to_feature(source_global) == 0) THEN
785 2729390 : cached_layout%global_to_feature(source_global) = row
786 : END IF
787 3480442 : static_base = nstatic_per_point*(ipt - 1)
788 13921768 : cached_layout%grid_coords(:, row) = global_static(static_base + 1:static_base + 3)
789 3480442 : cached_layout%grid_weights(row) = global_static(static_base + 4)
790 3480442 : cached_layout%atomic_grid_weights(row) = global_static(static_base + 5)
791 3480442 : source_local = source_global - point_displs(pe_index)
792 3480570 : IF (source_local >= 1 .AND. source_local <= nflat_local) THEN
793 1740221 : feature_local = cursor(source_local)
794 1740221 : cursor(source_local) = cursor(source_local) + 1
795 1740221 : cached_layout%local_feature_rows(feature_local) = row
796 1740221 : cached_layout%local_feature_points(feature_local) = source_local
797 : END IF
798 : END DO
799 :
800 2729518 : CPASSERT(ALL(cached_layout%global_to_feature > 0))
801 1740349 : CPASSERT(ALL(cached_layout%local_feature_rows > 0))
802 1740349 : CPASSERT(ALL(cached_layout%local_feature_points > 0))
803 3118 : DO k = bo(1, 3), bo(2, 3)
804 86436 : DO j = bo(1, 2), bo(2, 2)
805 1451003 : DO i = bo(1, 1), bo(2, 1)
806 1364695 : local_row = cached_layout%feature_index(i, j, k)
807 : cached_layout%feature_index(i, j, k) = &
808 1448013 : cached_layout%local_feature_rows(cached_layout%local_feature_offsets(local_row))
809 : END DO
810 : END DO
811 : END DO
812 1740349 : DO feature_local = 1, nfeature_local
813 : cached_layout%local_feature_indices(feature_local) = &
814 1740349 : INT(cached_layout%local_feature_rows(feature_local) - 1, KIND=int_8)
815 : END DO
816 128 : CALL timestop(phase_handle)
817 128 : CALL timeset("skala_gpw_layout_chunk_routes", phase_handle)
818 : CALL build_atom_chunk_routes(cached_layout, cached_layout%local_feature_rows, &
819 128 : pw_grid%para%group)
820 128 : CALL build_atom_chunk_layout(cached_layout)
821 128 : CALL timestop(phase_handle)
822 :
823 128 : cached_layout%natom = natom
824 128 : cached_layout%nflat = nflat
825 128 : cached_layout%nflat_local = nflat_local
826 128 : cached_layout%npoint = npoint
827 128 : cached_layout%nproc = nproc
828 128 : cached_layout%atom_partition = my_atom_partition
829 1280 : cached_layout%bo = bo
830 1280 : cached_layout%bounds = pw_grid%bounds
831 512 : cached_layout%npts = pw_grid%npts
832 128 : cached_layout%dvol = pw_grid%dvol
833 1664 : cached_layout%dh = pw_grid%dh
834 1664 : cached_layout%cell_hmat = cell%hmat
835 128 : cached_layout%weight_sum = weight_sum
836 128 : cached_layout%weight_sumsq = weight_sumsq
837 128 : cached_layout%has_weights = has_weights
838 128 : CALL timeset("skala_gpw_layout_tensors", phase_handle)
839 128 : CALL build_static_layout_tensors(cached_layout)
840 128 : CALL timestop(phase_handle)
841 128 : cached_layout%active = .TRUE.
842 :
843 0 : DEALLOCATE (atom_coords_pbc, atom_image_coords, atom_offset, atom_position, &
844 0 : chunk_atom_begin, chunk_atom_end, cursor, feature_counts, feature_displs, &
845 0 : global_owner, global_source_points, global_static, local_feature_counts_tmp, &
846 0 : distances, local_owner, local_source_global, local_source_points, &
847 0 : local_static, partition_weights, point_counts, point_displs, static_counts, &
848 128 : static_displs)
849 :
850 640 : END SUBROUTINE rebuild_layout_cache
851 :
852 : ! **************************************************************************************************
853 : !> \brief Build cached Torch tensors for static SKALA inputs.
854 : !> \param cache ...
855 : ! **************************************************************************************************
856 128 : SUBROUTINE build_static_layout_tensors(cache)
857 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
858 :
859 128 : CPASSERT(.NOT. cache%static_tensors_active)
860 :
861 128 : CALL torch_tensor_from_array(cache%grid_coords_t, cache%grid_coords)
862 128 : CALL torch_tensor_to_device_leaf(cache%grid_coords_t, .FALSE.)
863 128 : CALL torch_tensor_from_array(cache%grid_weights_t, cache%grid_weights)
864 128 : CALL torch_tensor_to_device_leaf(cache%grid_weights_t, .FALSE.)
865 128 : CALL torch_tensor_from_array(cache%atomic_grid_weights_t, cache%atomic_grid_weights)
866 128 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_weights_t, .FALSE.)
867 128 : CALL torch_tensor_from_array(cache%atomic_grid_sizes_t, cache%atomic_grid_sizes)
868 128 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_sizes_t, .FALSE.)
869 128 : CALL torch_tensor_from_array(cache%coarse_0_atomic_coords_t, cache%coarse_0_atomic_coords)
870 128 : CALL torch_tensor_to_device_leaf(cache%coarse_0_atomic_coords_t, .FALSE.)
871 : CALL torch_tensor_from_array(cache%atomic_grid_size_bound_shape_t, &
872 128 : cache%atomic_grid_size_bound_shape)
873 128 : CALL torch_tensor_to_device_leaf(cache%atomic_grid_size_bound_shape_t, .FALSE.)
874 128 : CALL torch_tensor_from_array(cache%local_feature_indices_t, cache%local_feature_indices)
875 128 : CALL torch_tensor_to_device_leaf(cache%local_feature_indices_t, .FALSE.)
876 :
877 128 : CALL torch_dict_create(cache%static_inputs)
878 128 : CALL torch_dict_insert(cache%static_inputs, "grid_coords", cache%grid_coords_t)
879 128 : CALL torch_dict_insert(cache%static_inputs, "grid_weights", cache%grid_weights_t)
880 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_weights", &
881 128 : cache%atomic_grid_weights_t)
882 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_sizes", &
883 128 : cache%atomic_grid_sizes_t)
884 : CALL torch_dict_insert(cache%static_inputs, "atomic_grid_size_bound_shape", &
885 128 : cache%atomic_grid_size_bound_shape_t)
886 128 : cache%static_tensors_active = .TRUE.
887 :
888 128 : IF (cache%chunk_feature_count > 0) THEN
889 128 : CPASSERT(.NOT. cache%chunk_static_tensors_active)
890 128 : CALL torch_tensor_from_array(cache%chunk_grid_coords_t, cache%chunk_grid_coords)
891 128 : CALL torch_tensor_to_device_leaf(cache%chunk_grid_coords_t, .FALSE.)
892 128 : CALL torch_tensor_from_array(cache%chunk_grid_weights_t, cache%chunk_grid_weights)
893 128 : CALL torch_tensor_to_device_leaf(cache%chunk_grid_weights_t, .FALSE.)
894 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_weights_t, &
895 128 : cache%chunk_atomic_grid_weights)
896 128 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_weights_t, .FALSE.)
897 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_sizes_t, &
898 128 : cache%chunk_atomic_grid_sizes)
899 128 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_sizes_t, .FALSE.)
900 : CALL torch_tensor_from_array(cache%chunk_coarse_0_atomic_coords_t, &
901 128 : cache%chunk_coarse_0_atomic_coords)
902 128 : CALL torch_tensor_to_device_leaf(cache%chunk_coarse_0_atomic_coords_t, .FALSE.)
903 : CALL torch_tensor_from_array(cache%chunk_atomic_grid_size_bound_shape_t, &
904 128 : cache%chunk_atomic_grid_size_bound_shape)
905 128 : CALL torch_tensor_to_device_leaf(cache%chunk_atomic_grid_size_bound_shape_t, .FALSE.)
906 128 : CALL torch_tensor_from_array(cache%chunk_feature_indices_t, cache%chunk_feature_indices)
907 128 : CALL torch_tensor_to_device_leaf(cache%chunk_feature_indices_t, .FALSE.)
908 :
909 128 : CALL torch_dict_create(cache%chunk_static_inputs)
910 : CALL torch_dict_insert(cache%chunk_static_inputs, "grid_coords", &
911 128 : cache%chunk_grid_coords_t)
912 : CALL torch_dict_insert(cache%chunk_static_inputs, "grid_weights", &
913 128 : cache%chunk_grid_weights_t)
914 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_weights", &
915 128 : cache%chunk_atomic_grid_weights_t)
916 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_sizes", &
917 128 : cache%chunk_atomic_grid_sizes_t)
918 : CALL torch_dict_insert(cache%chunk_static_inputs, "atomic_grid_size_bound_shape", &
919 128 : cache%chunk_atomic_grid_size_bound_shape_t)
920 128 : cache%chunk_static_tensors_active = .TRUE.
921 : END IF
922 :
923 128 : END SUBROUTINE build_static_layout_tensors
924 :
925 : ! **************************************************************************************************
926 : !> \brief Copy static cached layout arrays into a feature bundle.
927 : !> \param features ...
928 : !> \param needs_coordinate_array ...
929 : !> \param needs_grid_coordinate_array ...
930 : ! **************************************************************************************************
931 290 : SUBROUTINE copy_cached_layout(features, needs_coordinate_array, needs_grid_coordinate_array)
932 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
933 : LOGICAL, INTENT(IN) :: needs_coordinate_array, &
934 : needs_grid_coordinate_array
935 :
936 290 : CPASSERT(cached_layout%active)
937 :
938 0 : ALLOCATE (features%feature_index(LBOUND(cached_layout%feature_index, 1): &
939 : UBOUND(cached_layout%feature_index, 1), &
940 : LBOUND(cached_layout%feature_index, 2): &
941 : UBOUND(cached_layout%feature_index, 2), &
942 : LBOUND(cached_layout%feature_index, 3): &
943 1450 : UBOUND(cached_layout%feature_index, 3)))
944 870 : ALLOCATE (features%grid_weights(cached_layout%nflat))
945 0 : ALLOCATE (features%local_feature_counts(cached_layout%nflat_local), &
946 0 : features%local_feature_offsets(cached_layout%nflat_local + 1), &
947 2030 : features%local_feature_rows(SIZE(cached_layout%local_feature_rows)))
948 :
949 1983331 : features%feature_index(:, :, :) = cached_layout%feature_index
950 5390600 : features%grid_weights(:) = cached_layout%grid_weights
951 1843743 : features%local_feature_counts(:) = cached_layout%local_feature_counts
952 1844033 : features%local_feature_offsets(:) = cached_layout%local_feature_offsets
953 2695445 : features%local_feature_rows(:) = cached_layout%local_feature_rows
954 290 : features%nflat = cached_layout%nflat
955 290 : features%nflat_local = cached_layout%nflat_local
956 290 : features%chunk_feature_count = cached_layout%chunk_feature_count
957 290 : features%atom_partition = cached_layout%atom_partition
958 870 : ALLOCATE (features%atomic_grid_sizes(cached_layout%natom))
959 898 : features%atomic_grid_sizes(:) = cached_layout%atomic_grid_sizes
960 290 : IF (needs_grid_coordinate_array) THEN
961 180 : ALLOCATE (features%grid_coords(3, cached_layout%nflat))
962 120 : ALLOCATE (features%atomic_grid_weights(cached_layout%nflat))
963 4727740 : features%grid_coords(:, :) = cached_layout%grid_coords
964 1181980 : features%atomic_grid_weights(:) = cached_layout%atomic_grid_weights
965 : END IF
966 0 : ALLOCATE (features%chunk_grad_counts(cached_layout%nproc), &
967 0 : features%chunk_grad_displs(cached_layout%nproc), &
968 0 : features%route_grad_return_recv_counts(cached_layout%nproc), &
969 0 : features%route_grad_return_recv_displs(cached_layout%nproc), &
970 0 : features%route_grad_return_send_counts(cached_layout%nproc), &
971 0 : features%route_grad_return_send_displs(cached_layout%nproc), &
972 0 : features%route_point_recv_counts(cached_layout%nproc), &
973 0 : features%route_point_recv_displs(cached_layout%nproc), &
974 0 : features%route_point_send_counts(cached_layout%nproc), &
975 0 : features%route_point_send_displs(cached_layout%nproc), &
976 4060 : features%route_send_local_rows(SIZE(cached_layout%route_send_local_rows)))
977 870 : features%chunk_grad_counts(:) = cached_layout%chunk_grad_counts
978 870 : features%chunk_grad_displs(:) = cached_layout%chunk_grad_displs
979 870 : features%route_grad_return_recv_counts(:) = cached_layout%route_grad_return_recv_counts
980 870 : features%route_grad_return_recv_displs(:) = cached_layout%route_grad_return_recv_displs
981 870 : features%route_grad_return_send_counts(:) = cached_layout%route_grad_return_send_counts
982 870 : features%route_grad_return_send_displs(:) = cached_layout%route_grad_return_send_displs
983 870 : features%route_point_recv_counts(:) = cached_layout%route_point_recv_counts
984 870 : features%route_point_recv_displs(:) = cached_layout%route_point_recv_displs
985 870 : features%route_point_send_counts(:) = cached_layout%route_point_send_counts
986 870 : features%route_point_send_displs(:) = cached_layout%route_point_send_displs
987 2695445 : features%route_send_local_rows(:) = cached_layout%route_send_local_rows
988 290 : IF (needs_coordinate_array) THEN
989 180 : ALLOCATE (features%coarse_0_atomic_coords(3, cached_layout%natom))
990 540 : features%coarse_0_atomic_coords(:, :) = cached_layout%coarse_0_atomic_coords
991 : END IF
992 :
993 290 : END SUBROUTINE copy_cached_layout
994 :
995 : ! **************************************************************************************************
996 : !> \brief Split the atom-ordered feature rows into contiguous atom chunks.
997 : !> \param atomic_grid_sizes ...
998 : !> \param atom_offset ...
999 : !> \param nproc ...
1000 : !> \param chunk_atom_begin ...
1001 : !> \param chunk_atom_end ...
1002 : !> \param chunk_feature_counts ...
1003 : !> \param chunk_feature_displs ...
1004 : ! **************************************************************************************************
1005 128 : SUBROUTINE build_atom_chunks(atomic_grid_sizes, atom_offset, nproc, chunk_atom_begin, &
1006 128 : chunk_atom_end, chunk_feature_counts, chunk_feature_displs)
1007 : INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1008 : INTEGER, DIMENSION(:), INTENT(IN) :: atom_offset
1009 : INTEGER, INTENT(IN) :: nproc
1010 : INTEGER, DIMENSION(:), INTENT(OUT) :: chunk_atom_begin, chunk_atom_end, &
1011 : chunk_feature_counts, &
1012 : chunk_feature_displs
1013 :
1014 : INTEGER :: best_limit, count, displ, end_atom, lower_limit, max_end_atom, midpoint, natom, &
1015 : next_atom, next_count, pe, ranks_left, target_chunks, total_count, upper_limit
1016 :
1017 128 : natom = SIZE(atomic_grid_sizes)
1018 384 : chunk_atom_begin = natom + 1
1019 384 : chunk_atom_end = natom
1020 384 : chunk_feature_counts = 0
1021 384 : chunk_feature_displs = 0
1022 128 : IF (natom == 0) RETURN
1023 :
1024 128 : target_chunks = MIN(nproc, natom)
1025 128 : total_count = atom_offset(natom + 1) - 1
1026 412 : lower_limit = MAXVAL(INT(atomic_grid_sizes))
1027 128 : lower_limit = MAX(lower_limit, (total_count + target_chunks - 1)/target_chunks)
1028 128 : upper_limit = total_count
1029 128 : best_limit = upper_limit
1030 1722 : DO WHILE (lower_limit <= upper_limit)
1031 1594 : midpoint = (lower_limit + upper_limit)/2
1032 1722 : IF (atom_chunks_fit_limit(atomic_grid_sizes, midpoint, target_chunks)) THEN
1033 1474 : best_limit = midpoint
1034 1474 : upper_limit = midpoint - 1
1035 : ELSE
1036 120 : lower_limit = midpoint + 1
1037 : END IF
1038 : END DO
1039 :
1040 : displ = 0
1041 : next_atom = 1
1042 384 : DO pe = 1, nproc
1043 256 : chunk_feature_displs(pe) = displ
1044 256 : IF (pe > target_chunks .OR. next_atom > natom) CYCLE
1045 :
1046 256 : ranks_left = target_chunks - pe + 1
1047 256 : chunk_atom_begin(pe) = next_atom
1048 256 : max_end_atom = natom - ranks_left + 1
1049 256 : end_atom = next_atom
1050 256 : count = INT(atomic_grid_sizes(end_atom))
1051 284 : DO WHILE (end_atom < max_end_atom)
1052 38 : next_count = count + INT(atomic_grid_sizes(end_atom + 1))
1053 38 : IF (next_count > best_limit) EXIT
1054 : end_atom = end_atom + 1
1055 256 : count = next_count
1056 : END DO
1057 :
1058 256 : chunk_atom_end(pe) = end_atom
1059 256 : chunk_feature_counts(pe) = atom_offset(end_atom + 1) - atom_offset(next_atom)
1060 256 : displ = displ + chunk_feature_counts(pe)
1061 384 : next_atom = end_atom + 1
1062 : END DO
1063 :
1064 128 : CPASSERT(displ == atom_offset(natom + 1) - 1)
1065 :
1066 : END SUBROUTINE build_atom_chunks
1067 :
1068 : ! **************************************************************************************************
1069 : !> \brief Check if contiguous atom chunks can stay below a feature-count limit.
1070 : !> \param atomic_grid_sizes ...
1071 : !> \param limit ...
1072 : !> \param nchunks ...
1073 : !> \return ...
1074 : ! **************************************************************************************************
1075 1594 : FUNCTION atom_chunks_fit_limit(atomic_grid_sizes, limit, nchunks) RESULT(fits)
1076 : INTEGER(KIND=int_8), DIMENSION(:), INTENT(IN) :: atomic_grid_sizes
1077 : INTEGER, INTENT(IN) :: limit, nchunks
1078 : LOGICAL :: fits
1079 :
1080 : INTEGER :: atom_count, chunk_count, iatom, &
1081 : used_chunks
1082 :
1083 1594 : fits = .FALSE.
1084 1594 : IF (SIZE(atomic_grid_sizes) == 0) THEN
1085 1594 : fits = .TRUE.
1086 : RETURN
1087 : END IF
1088 :
1089 5202 : used_chunks = 1
1090 5202 : chunk_count = 0
1091 5202 : DO iatom = 1, SIZE(atomic_grid_sizes)
1092 3608 : atom_count = INT(atomic_grid_sizes(iatom))
1093 3608 : IF (atom_count > limit) RETURN
1094 5202 : IF (chunk_count + atom_count > limit) THEN
1095 1714 : used_chunks = used_chunks + 1
1096 1714 : chunk_count = atom_count
1097 : ELSE
1098 : chunk_count = chunk_count + atom_count
1099 : END IF
1100 : END DO
1101 1594 : fits = used_chunks <= nchunks
1102 :
1103 1594 : END FUNCTION atom_chunks_fit_limit
1104 :
1105 : ! **************************************************************************************************
1106 : !> \brief Return the MPI rank owning an atom-ordered feature row.
1107 : !> \param row ...
1108 : !> \param counts ...
1109 : !> \param displs ...
1110 : !> \return ...
1111 : ! **************************************************************************************************
1112 1740221 : FUNCTION feature_row_chunk_owner(row, counts, displs) RESULT(owner)
1113 : INTEGER, INTENT(IN) :: row
1114 : INTEGER, DIMENSION(:), INTENT(IN) :: counts, displs
1115 : INTEGER :: owner
1116 :
1117 : INTEGER :: pe
1118 :
1119 1740221 : owner = 0
1120 2569695 : DO pe = 1, SIZE(counts)
1121 2569695 : IF (row > displs(pe) .AND. row <= displs(pe) + counts(pe)) THEN
1122 1740221 : owner = pe
1123 : RETURN
1124 : END IF
1125 : END DO
1126 :
1127 : END FUNCTION feature_row_chunk_owner
1128 :
1129 : ! **************************************************************************************************
1130 : !> \brief Build zero-based displacement arrays from per-rank counts.
1131 : !> \param counts ...
1132 : !> \param displs ...
1133 : ! **************************************************************************************************
1134 512 : SUBROUTINE counts_to_displs(counts, displs)
1135 : INTEGER, DIMENSION(:), INTENT(IN) :: counts
1136 : INTEGER, DIMENSION(:), INTENT(OUT) :: displs
1137 :
1138 : INTEGER :: pe
1139 :
1140 512 : displs(1) = 0
1141 1024 : DO pe = 2, SIZE(counts)
1142 1024 : displs(pe) = displs(pe - 1) + counts(pe - 1)
1143 : END DO
1144 :
1145 512 : END SUBROUTINE counts_to_displs
1146 :
1147 : ! **************************************************************************************************
1148 : !> \brief Precompute all-to-all routing between local grid rows and atom chunks.
1149 : !> \param cache ...
1150 : !> \param local_to_global ...
1151 : !> \param group ...
1152 : ! **************************************************************************************************
1153 128 : SUBROUTINE build_atom_chunk_routes(cache, local_to_global, group)
1154 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1155 : INTEGER, DIMENSION(:), INTENT(IN) :: local_to_global
1156 :
1157 : CLASS(mp_comm_type), INTENT(IN) :: group
1158 :
1159 : INTEGER :: chunk_row, dest, local_feature, point_pos, row
1160 128 : INTEGER, ALLOCATABLE, DIMENSION(:) :: cursor, recv_meta, send_meta
1161 :
1162 0 : ALLOCATE (cache%route_local_dest(SIZE(local_to_global)), &
1163 0 : cache%route_send_local_rows(SIZE(local_to_global)), &
1164 0 : cache%chunk_return_positions(cache%chunk_feature_count), &
1165 1024 : cursor(SIZE(cache%route_point_send_counts)))
1166 384 : cache%route_point_send_counts = 0
1167 1740349 : cache%route_send_local_rows = 0
1168 1740349 : cache%chunk_return_positions = 0
1169 1740349 : DO local_feature = 1, SIZE(local_to_global)
1170 : dest = feature_row_chunk_owner(local_to_global(local_feature), &
1171 : cache%chunk_feature_counts, &
1172 1740221 : cache%chunk_feature_displs)
1173 1740221 : CPASSERT(dest > 0)
1174 1740221 : cache%route_local_dest(local_feature) = dest
1175 1740349 : cache%route_point_send_counts(dest) = cache%route_point_send_counts(dest) + 1
1176 : END DO
1177 128 : CALL counts_to_displs(cache%route_point_send_counts, cache%route_point_send_displs)
1178 384 : cursor(:) = cache%route_point_send_displs + 1
1179 1740349 : DO local_feature = 1, SIZE(local_to_global)
1180 1740221 : dest = cache%route_local_dest(local_feature)
1181 1740221 : point_pos = cursor(dest)
1182 1740221 : cursor(dest) = cursor(dest) + 1
1183 1740349 : cache%route_send_local_rows(point_pos) = cache%local_feature_points(local_feature)
1184 : END DO
1185 128 : CALL group%alltoall(cache%route_point_send_counts, cache%route_point_recv_counts, 1)
1186 128 : CALL counts_to_displs(cache%route_point_recv_counts, cache%route_point_recv_displs)
1187 :
1188 512 : ALLOCATE (send_meta(SIZE(local_to_global)), recv_meta(cache%chunk_feature_count))
1189 384 : cursor(:) = cache%route_point_send_displs + 1
1190 1740349 : DO local_feature = 1, SIZE(local_to_global)
1191 1740221 : dest = cache%route_local_dest(local_feature)
1192 1740221 : point_pos = cursor(dest)
1193 1740221 : cursor(dest) = cursor(dest) + 1
1194 1740349 : send_meta(point_pos) = local_to_global(local_feature)
1195 : END DO
1196 : CALL group%alltoall(send_meta, cache%route_point_send_counts, &
1197 : cache%route_point_send_displs, recv_meta, &
1198 : cache%route_point_recv_counts, &
1199 128 : cache%route_point_recv_displs)
1200 1740349 : DO point_pos = 1, cache%chunk_feature_count
1201 1740221 : row = recv_meta(point_pos)
1202 1740221 : chunk_row = row - cache%chunk_feature_begin + 1
1203 1740221 : CPASSERT(chunk_row >= 1 .AND. chunk_row <= cache%chunk_feature_count)
1204 1740349 : cache%chunk_return_positions(chunk_row) = point_pos
1205 : END DO
1206 :
1207 384 : cache%route_grad_return_send_counts(:) = ngrad_per_point*cache%route_point_recv_counts
1208 384 : cache%route_grad_return_send_displs(:) = ngrad_per_point*cache%route_point_recv_displs
1209 384 : cache%route_grad_return_recv_counts(:) = ngrad_per_point*cache%route_point_send_counts
1210 384 : cache%route_grad_return_recv_displs(:) = ngrad_per_point*cache%route_point_send_displs
1211 :
1212 384 : CPASSERT(SUM(cache%route_point_send_counts) == SIZE(local_to_global))
1213 384 : CPASSERT(SUM(cache%route_point_recv_counts) == cache%chunk_feature_count)
1214 1740349 : CPASSERT(ALL(cache%route_send_local_rows > 0))
1215 1740349 : CPASSERT(ALL(cache%chunk_return_positions > 0))
1216 :
1217 128 : DEALLOCATE (cursor, recv_meta, send_meta)
1218 :
1219 128 : END SUBROUTINE build_atom_chunk_routes
1220 :
1221 : ! **************************************************************************************************
1222 : !> \brief Materialize the current rank's atom chunk static layout.
1223 : !> \param cache ...
1224 : ! **************************************************************************************************
1225 128 : SUBROUTINE build_atom_chunk_layout(cache)
1226 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1227 :
1228 : INTEGER :: irow, max_grid_size, row_begin, row_end
1229 :
1230 128 : IF (cache%chunk_feature_count <= 0 .OR. cache%chunk_natom <= 0) RETURN
1231 :
1232 128 : row_begin = cache%chunk_feature_begin
1233 128 : row_end = row_begin + cache%chunk_feature_count - 1
1234 0 : ALLOCATE (cache%chunk_grid_coords(3, cache%chunk_feature_count), &
1235 0 : cache%chunk_grid_weights(cache%chunk_feature_count), &
1236 0 : cache%chunk_atomic_grid_weights(cache%chunk_feature_count), &
1237 0 : cache%chunk_atomic_grid_sizes(cache%chunk_natom), &
1238 0 : cache%chunk_coarse_0_atomic_coords(3, cache%chunk_natom), &
1239 1408 : cache%chunk_feature_indices(cache%chunk_feature_count))
1240 6961012 : cache%chunk_grid_coords(:, :) = cache%grid_coords(:, row_begin:row_end)
1241 1740349 : cache%chunk_grid_weights(:) = cache%grid_weights(row_begin:row_end)
1242 1740349 : cache%chunk_atomic_grid_weights(:) = cache%atomic_grid_weights(row_begin:row_end)
1243 : cache%chunk_atomic_grid_sizes(:) = &
1244 270 : cache%atomic_grid_sizes(cache%chunk_atom_begin:cache%chunk_atom_end)
1245 : cache%chunk_coarse_0_atomic_coords(:, :) = &
1246 696 : cache%coarse_0_atomic_coords(:, cache%chunk_atom_begin:cache%chunk_atom_end)
1247 :
1248 270 : max_grid_size = MAXVAL(INT(cache%chunk_atomic_grid_sizes))
1249 256 : ALLOCATE (cache%chunk_atomic_grid_size_bound_shape(0, max_grid_size))
1250 1519358 : cache%chunk_atomic_grid_size_bound_shape = 0_int_8
1251 1740349 : DO irow = 1, cache%chunk_feature_count
1252 1740349 : cache%chunk_feature_indices(irow) = INT(irow - 1, KIND=int_8)
1253 : END DO
1254 :
1255 : END SUBROUTINE build_atom_chunk_layout
1256 :
1257 : ! **************************************************************************************************
1258 : !> \brief Send local dynamic feature rows to their atom-chunk owner ranks.
1259 : !> \param features ...
1260 : !> \param local_dynamic ...
1261 : !> \param group ...
1262 : !> \param collapse_spin_dynamics ...
1263 : ! **************************************************************************************************
1264 8 : SUBROUTINE route_atom_chunk_dynamics(features, local_dynamic, group, collapse_spin_dynamics)
1265 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1266 : REAL(KIND=dp), DIMENSION(:), INTENT(IN) :: local_dynamic
1267 :
1268 : CLASS(mp_comm_type), INTENT(IN) :: group
1269 : LOGICAL, INTENT(IN) :: collapse_spin_dynamics
1270 :
1271 : INTEGER :: chunk_row, dest, dyn_base, local_feature, local_row, &
1272 : ndynamic_route_per_point, nrecv, nsend, &
1273 : point_pos, src_base
1274 8 : INTEGER, ALLOCATABLE, DIMENSION(:) :: cursor, recv_counts, recv_displs, &
1275 : send_counts, send_displs
1276 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: recv_dynamic, send_dynamic
1277 :
1278 8 : nsend = SIZE(cached_layout%route_local_dest)
1279 24 : nrecv = SUM(cached_layout%route_point_recv_counts)
1280 8 : CPASSERT(nsend == SIZE(cached_layout%local_feature_rows))
1281 8 : CPASSERT(nrecv == cached_layout%chunk_feature_count)
1282 8 : ndynamic_route_per_point = ndynamic_per_point
1283 8 : IF (collapse_spin_dynamics) ndynamic_route_per_point = nrks_dynamic_per_point
1284 :
1285 : ALLOCATE (send_dynamic(MAX(1, ndynamic_route_per_point*nsend)), &
1286 : recv_dynamic(MAX(1, ndynamic_route_per_point*nrecv)), &
1287 : cursor(cached_layout%nproc), send_counts(cached_layout%nproc), &
1288 : send_displs(cached_layout%nproc), recv_counts(cached_layout%nproc), &
1289 88 : recv_displs(cached_layout%nproc))
1290 24 : send_counts(:) = ndynamic_route_per_point*cached_layout%route_point_send_counts
1291 24 : send_displs(:) = ndynamic_route_per_point*cached_layout%route_point_send_displs
1292 24 : recv_counts(:) = ndynamic_route_per_point*cached_layout%route_point_recv_counts
1293 24 : recv_displs(:) = ndynamic_route_per_point*cached_layout%route_point_recv_displs
1294 24 : cursor(:) = cached_layout%route_point_send_displs + 1
1295 146742 : DO local_feature = 1, nsend
1296 146734 : dest = cached_layout%route_local_dest(local_feature)
1297 146734 : point_pos = cursor(dest)
1298 146734 : cursor(dest) = cursor(dest) + 1
1299 146734 : dyn_base = ndynamic_route_per_point*(point_pos - 1)
1300 146734 : local_row = cached_layout%local_feature_points(local_feature)
1301 146734 : src_base = ndynamic_route_per_point*(local_row - 1)
1302 : send_dynamic(dyn_base + 1:dyn_base + ndynamic_route_per_point) = &
1303 880412 : local_dynamic(src_base + 1:src_base + ndynamic_route_per_point)
1304 : END DO
1305 :
1306 : CALL group%alltoall(send_dynamic, send_counts, send_displs, recv_dynamic, recv_counts, &
1307 8 : recv_displs)
1308 :
1309 8 : features%uses_collapsed_rks_dynamic = collapse_spin_dynamics
1310 8 : IF (cached_layout%chunk_feature_count > 0) THEN
1311 8 : IF (collapse_spin_dynamics) THEN
1312 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 1), &
1313 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 1), &
1314 0 : features%chunk_kin(cached_layout%chunk_feature_count, 1), &
1315 64 : features%chunk_return_positions(cached_layout%chunk_feature_count))
1316 : ELSE
1317 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1318 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1319 0 : features%chunk_kin(cached_layout%chunk_feature_count, 2), &
1320 0 : features%chunk_return_positions(cached_layout%chunk_feature_count))
1321 : END IF
1322 146742 : features%chunk_return_positions(:) = cached_layout%chunk_return_positions
1323 :
1324 146742 : DO chunk_row = 1, cached_layout%chunk_feature_count
1325 146734 : point_pos = cached_layout%chunk_return_positions(chunk_row)
1326 146734 : CPASSERT(point_pos >= 1 .AND. point_pos <= cached_layout%chunk_feature_count)
1327 146734 : dyn_base = ndynamic_route_per_point*(point_pos - 1)
1328 146742 : IF (collapse_spin_dynamics) THEN
1329 146734 : features%chunk_density(chunk_row, 1) = recv_dynamic(dyn_base + 1)
1330 146734 : features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 2)
1331 146734 : features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 3)
1332 146734 : features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 4)
1333 146734 : features%chunk_kin(chunk_row, 1) = recv_dynamic(dyn_base + 5)
1334 : ELSE
1335 0 : features%chunk_density(chunk_row, :) = recv_dynamic(dyn_base + 1:dyn_base + 2)
1336 0 : features%chunk_grad(chunk_row, 1, 1) = recv_dynamic(dyn_base + 3)
1337 0 : features%chunk_grad(chunk_row, 2, 1) = recv_dynamic(dyn_base + 4)
1338 0 : features%chunk_grad(chunk_row, 3, 1) = recv_dynamic(dyn_base + 5)
1339 0 : features%chunk_grad(chunk_row, 1, 2) = recv_dynamic(dyn_base + 6)
1340 0 : features%chunk_grad(chunk_row, 2, 2) = recv_dynamic(dyn_base + 7)
1341 0 : features%chunk_grad(chunk_row, 3, 2) = recv_dynamic(dyn_base + 8)
1342 0 : features%chunk_kin(chunk_row, :) = recv_dynamic(dyn_base + 9:dyn_base + 10)
1343 : END IF
1344 : END DO
1345 146742 : CPASSERT(ALL(features%chunk_return_positions > 0))
1346 : END IF
1347 :
1348 0 : DEALLOCATE (cursor, recv_counts, recv_displs, recv_dynamic, send_counts, send_displs, &
1349 8 : send_dynamic)
1350 :
1351 8 : END SUBROUTINE route_atom_chunk_dynamics
1352 :
1353 : ! **************************************************************************************************
1354 : !> \brief Extract the current rank's atom chunk from the global dynamic feature arrays.
1355 : !> \param features ...
1356 : ! **************************************************************************************************
1357 0 : SUBROUTINE extract_atom_chunk_dynamics(features)
1358 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1359 :
1360 : INTEGER :: row_begin, row_end
1361 :
1362 0 : CPASSERT(cached_layout%chunk_feature_count > 0)
1363 0 : row_begin = cached_layout%chunk_feature_begin
1364 0 : row_end = row_begin + cached_layout%chunk_feature_count - 1
1365 0 : ALLOCATE (features%chunk_density(cached_layout%chunk_feature_count, 2), &
1366 0 : features%chunk_grad(cached_layout%chunk_feature_count, 3, 2), &
1367 0 : features%chunk_kin(cached_layout%chunk_feature_count, 2))
1368 0 : features%chunk_density(:, :) = features%density(row_begin:row_end, :)
1369 0 : features%chunk_grad(:, :, :) = features%grad(row_begin:row_end, :, :)
1370 0 : features%chunk_kin(:, :) = features%kin(row_begin:row_end, :)
1371 :
1372 0 : END SUBROUTINE extract_atom_chunk_dynamics
1373 :
1374 : ! **************************************************************************************************
1375 : !> \brief Compute a local signature for optional integration weights.
1376 : !> \param weights ...
1377 : !> \param has_weights ...
1378 : !> \param weight_sum ...
1379 : !> \param weight_sumsq ...
1380 : ! **************************************************************************************************
1381 290 : SUBROUTINE weights_signature(weights, has_weights, weight_sum, weight_sumsq)
1382 : TYPE(pw_r3d_rs_type), OPTIONAL, POINTER :: weights
1383 : LOGICAL, INTENT(OUT) :: has_weights
1384 : REAL(KIND=dp), INTENT(OUT) :: weight_sum, weight_sumsq
1385 :
1386 290 : has_weights = .FALSE.
1387 290 : weight_sum = 0.0_dp
1388 290 : weight_sumsq = 0.0_dp
1389 290 : IF (PRESENT(weights)) THEN
1390 290 : IF (ASSOCIATED(weights)) THEN
1391 0 : has_weights = .TRUE.
1392 0 : weight_sum = SUM(weights%array)
1393 0 : weight_sumsq = SUM(weights%array*weights%array)
1394 : END IF
1395 : END IF
1396 :
1397 290 : END SUBROUTINE weights_signature
1398 :
1399 : ! **************************************************************************************************
1400 : !> \brief Release cached layout arrays.
1401 : !> \param cache ...
1402 : ! **************************************************************************************************
1403 128 : SUBROUTINE release_layout_cache(cache)
1404 : TYPE(skala_gpw_layout_cache_type), INTENT(INOUT) :: cache
1405 :
1406 128 : IF (cache%inputs_active) THEN
1407 40 : CALL torch_dict_release(cache%inputs)
1408 40 : cache%inputs_active = .FALSE.
1409 : END IF
1410 :
1411 128 : IF (cache%chunk_inputs_active) THEN
1412 0 : CALL torch_dict_release(cache%chunk_inputs)
1413 0 : cache%chunk_inputs_active = .FALSE.
1414 : END IF
1415 :
1416 128 : IF (cache%dynamic_tensors_active) THEN
1417 40 : CALL torch_tensor_release(cache%density_t)
1418 40 : CALL torch_tensor_release(cache%grad_t)
1419 40 : CALL torch_tensor_release(cache%kin_t)
1420 40 : cache%dynamic_tensors_active = .FALSE.
1421 : END IF
1422 :
1423 128 : IF (cache%chunk_dynamic_tensors_active) THEN
1424 0 : IF (cache%chunk_dynamic_input_views_active) THEN
1425 0 : CALL torch_tensor_release(cache%chunk_density_input_t)
1426 0 : CALL torch_tensor_release(cache%chunk_grad_input_t)
1427 0 : CALL torch_tensor_release(cache%chunk_kin_input_t)
1428 0 : cache%chunk_dynamic_input_views_active = .FALSE.
1429 : END IF
1430 0 : CALL torch_tensor_release(cache%chunk_density_t)
1431 0 : CALL torch_tensor_release(cache%chunk_grad_t)
1432 0 : CALL torch_tensor_release(cache%chunk_kin_t)
1433 0 : cache%chunk_dynamic_tensors_active = .FALSE.
1434 : END IF
1435 :
1436 128 : IF (cache%static_tensors_active) THEN
1437 40 : CALL torch_tensor_release(cache%grid_coords_t)
1438 40 : CALL torch_tensor_release(cache%grid_weights_t)
1439 40 : CALL torch_tensor_release(cache%atomic_grid_weights_t)
1440 40 : CALL torch_tensor_release(cache%atomic_grid_sizes_t)
1441 40 : CALL torch_tensor_release(cache%coarse_0_atomic_coords_t)
1442 40 : CALL torch_tensor_release(cache%atomic_grid_size_bound_shape_t)
1443 40 : CALL torch_tensor_release(cache%local_feature_indices_t)
1444 40 : CALL torch_dict_release(cache%static_inputs)
1445 40 : cache%static_tensors_active = .FALSE.
1446 : END IF
1447 :
1448 128 : IF (cache%chunk_static_tensors_active) THEN
1449 40 : CALL torch_tensor_release(cache%chunk_grid_coords_t)
1450 40 : CALL torch_tensor_release(cache%chunk_grid_weights_t)
1451 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_weights_t)
1452 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_sizes_t)
1453 40 : CALL torch_tensor_release(cache%chunk_coarse_0_atomic_coords_t)
1454 40 : CALL torch_tensor_release(cache%chunk_atomic_grid_size_bound_shape_t)
1455 40 : CALL torch_tensor_release(cache%chunk_feature_indices_t)
1456 40 : CALL torch_dict_release(cache%chunk_static_inputs)
1457 : cache%chunk_static_tensors_active = .FALSE.
1458 : END IF
1459 :
1460 128 : IF (ALLOCATED(cache%chunk_feature_counts)) DEALLOCATE (cache%chunk_feature_counts)
1461 128 : IF (ALLOCATED(cache%chunk_feature_displs)) DEALLOCATE (cache%chunk_feature_displs)
1462 128 : IF (ALLOCATED(cache%chunk_grad_counts)) DEALLOCATE (cache%chunk_grad_counts)
1463 128 : IF (ALLOCATED(cache%chunk_grad_displs)) DEALLOCATE (cache%chunk_grad_displs)
1464 128 : IF (ALLOCATED(cache%route_grad_return_recv_counts)) THEN
1465 40 : DEALLOCATE (cache%route_grad_return_recv_counts)
1466 : END IF
1467 128 : IF (ALLOCATED(cache%route_grad_return_recv_displs)) THEN
1468 40 : DEALLOCATE (cache%route_grad_return_recv_displs)
1469 : END IF
1470 128 : IF (ALLOCATED(cache%route_grad_return_send_counts)) THEN
1471 40 : DEALLOCATE (cache%route_grad_return_send_counts)
1472 : END IF
1473 128 : IF (ALLOCATED(cache%route_grad_return_send_displs)) THEN
1474 40 : DEALLOCATE (cache%route_grad_return_send_displs)
1475 : END IF
1476 128 : IF (ALLOCATED(cache%route_local_dest)) DEALLOCATE (cache%route_local_dest)
1477 128 : IF (ALLOCATED(cache%chunk_return_positions)) DEALLOCATE (cache%chunk_return_positions)
1478 128 : IF (ALLOCATED(cache%route_point_recv_counts)) DEALLOCATE (cache%route_point_recv_counts)
1479 128 : IF (ALLOCATED(cache%route_point_recv_displs)) DEALLOCATE (cache%route_point_recv_displs)
1480 128 : IF (ALLOCATED(cache%route_point_send_counts)) DEALLOCATE (cache%route_point_send_counts)
1481 128 : IF (ALLOCATED(cache%route_point_send_displs)) DEALLOCATE (cache%route_point_send_displs)
1482 128 : IF (ALLOCATED(cache%route_send_local_rows)) DEALLOCATE (cache%route_send_local_rows)
1483 128 : IF (ALLOCATED(cache%dynamic_counts)) DEALLOCATE (cache%dynamic_counts)
1484 128 : IF (ALLOCATED(cache%dynamic_displs)) DEALLOCATE (cache%dynamic_displs)
1485 128 : IF (ALLOCATED(cache%feature_counts)) DEALLOCATE (cache%feature_counts)
1486 128 : IF (ALLOCATED(cache%feature_displs)) DEALLOCATE (cache%feature_displs)
1487 128 : IF (ALLOCATED(cache%feature_source_points)) DEALLOCATE (cache%feature_source_points)
1488 128 : IF (ALLOCATED(cache%global_to_feature)) DEALLOCATE (cache%global_to_feature)
1489 128 : IF (ALLOCATED(cache%feature_index)) DEALLOCATE (cache%feature_index)
1490 128 : IF (ALLOCATED(cache%atomic_grid_sizes)) DEALLOCATE (cache%atomic_grid_sizes)
1491 128 : IF (ALLOCATED(cache%chunk_atomic_grid_sizes)) DEALLOCATE (cache%chunk_atomic_grid_sizes)
1492 128 : IF (ALLOCATED(cache%chunk_feature_indices)) DEALLOCATE (cache%chunk_feature_indices)
1493 128 : IF (ALLOCATED(cache%local_feature_counts)) DEALLOCATE (cache%local_feature_counts)
1494 128 : IF (ALLOCATED(cache%local_feature_indices)) DEALLOCATE (cache%local_feature_indices)
1495 128 : IF (ALLOCATED(cache%local_feature_offsets)) DEALLOCATE (cache%local_feature_offsets)
1496 128 : IF (ALLOCATED(cache%local_feature_points)) DEALLOCATE (cache%local_feature_points)
1497 128 : IF (ALLOCATED(cache%local_feature_rows)) DEALLOCATE (cache%local_feature_rows)
1498 128 : IF (ALLOCATED(cache%atomic_grid_size_bound_shape)) THEN
1499 40 : DEALLOCATE (cache%atomic_grid_size_bound_shape)
1500 : END IF
1501 128 : IF (ALLOCATED(cache%chunk_atomic_grid_size_bound_shape)) THEN
1502 40 : DEALLOCATE (cache%chunk_atomic_grid_size_bound_shape)
1503 : END IF
1504 128 : IF (ALLOCATED(cache%atomic_grid_weights)) DEALLOCATE (cache%atomic_grid_weights)
1505 128 : IF (ALLOCATED(cache%chunk_atomic_grid_weights)) DEALLOCATE (cache%chunk_atomic_grid_weights)
1506 128 : IF (ALLOCATED(cache%chunk_grid_weights)) DEALLOCATE (cache%chunk_grid_weights)
1507 128 : IF (ALLOCATED(cache%grid_weights)) DEALLOCATE (cache%grid_weights)
1508 128 : IF (ALLOCATED(cache%atom_coords)) DEALLOCATE (cache%atom_coords)
1509 128 : IF (ALLOCATED(cache%chunk_coarse_0_atomic_coords)) THEN
1510 40 : DEALLOCATE (cache%chunk_coarse_0_atomic_coords)
1511 : END IF
1512 128 : IF (ALLOCATED(cache%coarse_0_atomic_coords)) DEALLOCATE (cache%coarse_0_atomic_coords)
1513 128 : IF (ALLOCATED(cache%chunk_grid_coords)) DEALLOCATE (cache%chunk_grid_coords)
1514 128 : IF (ALLOCATED(cache%grid_coords)) DEALLOCATE (cache%grid_coords)
1515 :
1516 128 : cache%chunk_atom_begin = 1
1517 128 : cache%chunk_atom_end = 0
1518 128 : cache%chunk_feature_begin = 1
1519 128 : cache%chunk_feature_count = 0
1520 128 : cache%chunk_natom = 0
1521 128 : cache%natom = 0
1522 128 : cache%nflat = 0
1523 128 : cache%nflat_local = 0
1524 128 : cache%npoint = 0
1525 128 : cache%nproc = 0
1526 128 : cache%atom_partition = skala_gpw_atom_partition_hard
1527 1280 : cache%bo = 0
1528 1280 : cache%bounds = 0
1529 512 : cache%npts = 0
1530 128 : cache%dvol = 0.0_dp
1531 128 : cache%weight_sum = 0.0_dp
1532 128 : cache%weight_sumsq = 0.0_dp
1533 1664 : cache%cell_hmat = 0.0_dp
1534 1664 : cache%dh = 0.0_dp
1535 128 : cache%active = .FALSE.
1536 128 : cache%has_weights = .FALSE.
1537 128 : cache%chunk_dynamic_tensors_active = .FALSE.
1538 128 : cache%chunk_dynamic_input_views_active = .FALSE.
1539 128 : cache%chunk_inputs_active = .FALSE.
1540 128 : cache%chunk_inputs_use_collapsed_rks = .FALSE.
1541 128 : cache%chunk_static_tensors_active = .FALSE.
1542 128 : cache%dynamic_tensors_active = .FALSE.
1543 128 : cache%inputs_active = .FALSE.
1544 128 : cache%static_tensors_active = .FALSE.
1545 :
1546 128 : END SUBROUTINE release_layout_cache
1547 :
1548 : ! **************************************************************************************************
1549 : !> \brief Release Torch objects and backing arrays owned by a feature bundle.
1550 : !> \param features ...
1551 : ! **************************************************************************************************
1552 588 : SUBROUTINE skala_gpw_feature_release(features)
1553 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1554 :
1555 588 : IF (features%active) THEN
1556 294 : IF (features%owns_dynamic_tensors) THEN
1557 4 : IF (features%uses_collapsed_rks_dynamic) THEN
1558 4 : CALL torch_tensor_release(features%density_input_t)
1559 4 : CALL torch_tensor_release(features%grad_input_t)
1560 4 : CALL torch_tensor_release(features%kin_input_t)
1561 : END IF
1562 4 : CALL torch_tensor_release(features%density_t)
1563 4 : CALL torch_tensor_release(features%grad_t)
1564 4 : CALL torch_tensor_release(features%kin_t)
1565 : END IF
1566 294 : IF (features%owns_static_tensors) THEN
1567 4 : CALL torch_tensor_release(features%grid_coords_t)
1568 4 : CALL torch_tensor_release(features%grid_weights_t)
1569 4 : CALL torch_tensor_release(features%atomic_grid_weights_t)
1570 4 : CALL torch_tensor_release(features%atomic_grid_sizes_t)
1571 4 : CALL torch_tensor_release(features%atomic_grid_size_bound_shape_t)
1572 : END IF
1573 294 : IF (features%owns_grid_coordinate_tensor) THEN
1574 50 : CALL torch_tensor_release(features%grid_coords_t)
1575 : END IF
1576 294 : IF (features%owns_weight_tensors) THEN
1577 60 : CALL torch_tensor_release(features%grid_weights_t)
1578 60 : CALL torch_tensor_release(features%atomic_grid_weights_t)
1579 : END IF
1580 294 : IF (features%owns_static_tensors .OR. features%owns_coordinate_tensor) THEN
1581 64 : CALL torch_tensor_release(features%coarse_0_atomic_coords_t)
1582 : END IF
1583 294 : IF (features%owns_inputs) CALL torch_dict_release(features%inputs)
1584 294 : features%active = .FALSE.
1585 294 : features%owns_coordinate_tensor = .FALSE.
1586 294 : features%owns_grid_coordinate_tensor = .FALSE.
1587 294 : features%owns_weight_tensors = .FALSE.
1588 294 : features%owns_dynamic_tensors = .TRUE.
1589 294 : features%owns_inputs = .TRUE.
1590 294 : features%owns_static_tensors = .TRUE.
1591 : features%uses_atom_chunk_routing = .FALSE.
1592 294 : features%uses_atom_chunks = .FALSE.
1593 : features%uses_collapsed_rks_dynamic = .FALSE.
1594 : END IF
1595 :
1596 588 : IF (ALLOCATED(features%chunk_density)) DEALLOCATE (features%chunk_density)
1597 588 : IF (ALLOCATED(features%chunk_grad)) DEALLOCATE (features%chunk_grad)
1598 588 : IF (ALLOCATED(features%chunk_kin)) DEALLOCATE (features%chunk_kin)
1599 588 : IF (ALLOCATED(features%density)) DEALLOCATE (features%density)
1600 588 : IF (ALLOCATED(features%grad)) DEALLOCATE (features%grad)
1601 588 : IF (ALLOCATED(features%kin)) DEALLOCATE (features%kin)
1602 588 : IF (ALLOCATED(features%chunk_grad_counts)) DEALLOCATE (features%chunk_grad_counts)
1603 588 : IF (ALLOCATED(features%chunk_grad_displs)) DEALLOCATE (features%chunk_grad_displs)
1604 588 : IF (ALLOCATED(features%chunk_return_positions)) DEALLOCATE (features%chunk_return_positions)
1605 588 : IF (ALLOCATED(features%route_grad_return_recv_counts)) THEN
1606 290 : DEALLOCATE (features%route_grad_return_recv_counts)
1607 : END IF
1608 588 : IF (ALLOCATED(features%route_grad_return_recv_displs)) THEN
1609 290 : DEALLOCATE (features%route_grad_return_recv_displs)
1610 : END IF
1611 588 : IF (ALLOCATED(features%route_grad_return_send_counts)) THEN
1612 290 : DEALLOCATE (features%route_grad_return_send_counts)
1613 : END IF
1614 588 : IF (ALLOCATED(features%route_grad_return_send_displs)) THEN
1615 290 : DEALLOCATE (features%route_grad_return_send_displs)
1616 : END IF
1617 588 : IF (ALLOCATED(features%route_point_recv_counts)) THEN
1618 290 : DEALLOCATE (features%route_point_recv_counts)
1619 : END IF
1620 588 : IF (ALLOCATED(features%route_point_recv_displs)) THEN
1621 290 : DEALLOCATE (features%route_point_recv_displs)
1622 : END IF
1623 588 : IF (ALLOCATED(features%route_point_send_counts)) THEN
1624 290 : DEALLOCATE (features%route_point_send_counts)
1625 : END IF
1626 588 : IF (ALLOCATED(features%route_point_send_displs)) THEN
1627 290 : DEALLOCATE (features%route_point_send_displs)
1628 : END IF
1629 588 : IF (ALLOCATED(features%route_send_local_rows)) DEALLOCATE (features%route_send_local_rows)
1630 588 : IF (ALLOCATED(features%feature_index)) DEALLOCATE (features%feature_index)
1631 588 : IF (ALLOCATED(features%local_feature_counts)) DEALLOCATE (features%local_feature_counts)
1632 588 : IF (ALLOCATED(features%local_feature_offsets)) DEALLOCATE (features%local_feature_offsets)
1633 588 : IF (ALLOCATED(features%local_feature_rows)) DEALLOCATE (features%local_feature_rows)
1634 588 : IF (ALLOCATED(features%grid_coords)) DEALLOCATE (features%grid_coords)
1635 588 : IF (ALLOCATED(features%grid_weights)) DEALLOCATE (features%grid_weights)
1636 588 : IF (ALLOCATED(features%atomic_grid_weights)) DEALLOCATE (features%atomic_grid_weights)
1637 588 : IF (ALLOCATED(features%atomic_grid_sizes)) DEALLOCATE (features%atomic_grid_sizes)
1638 588 : IF (ALLOCATED(features%coarse_0_atomic_coords)) DEALLOCATE (features%coarse_0_atomic_coords)
1639 588 : IF (ALLOCATED(features%atomic_grid_size_bound_shape)) THEN
1640 4 : DEALLOCATE (features%atomic_grid_size_bound_shape)
1641 : END IF
1642 588 : features%chunk_feature_count = 0
1643 588 : features%nflat = 0
1644 588 : features%nflat_local = 0
1645 588 : features%atom_partition = skala_gpw_atom_partition_hard
1646 588 : features%uses_atom_chunk_routing = .FALSE.
1647 588 : features%uses_collapsed_rks_dynamic = .FALSE.
1648 :
1649 588 : END SUBROUTINE skala_gpw_feature_release
1650 :
1651 : ! **************************************************************************************************
1652 : !> \brief Return how many atom-contiguous subchunks the cached rank chunk needs.
1653 : !> \param max_rows ...
1654 : !> \return ...
1655 : ! **************************************************************************************************
1656 10 : FUNCTION skala_gpw_atom_subchunk_count(max_rows) RESULT(nsubchunks)
1657 : INTEGER, INTENT(IN) :: max_rows
1658 : INTEGER :: nsubchunks
1659 :
1660 : INTEGER :: atom_rows, iatom, rows
1661 :
1662 10 : nsubchunks = 0
1663 10 : IF (.NOT. cached_layout%active) RETURN
1664 10 : IF (cached_layout%chunk_natom <= 0) RETURN
1665 10 : IF (max_rows <= 0) THEN
1666 10 : nsubchunks = 1
1667 : RETURN
1668 : END IF
1669 :
1670 : rows = 0
1671 24 : DO iatom = 1, cached_layout%chunk_natom
1672 14 : atom_rows = INT(cached_layout%chunk_atomic_grid_sizes(iatom))
1673 14 : IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1674 4 : nsubchunks = nsubchunks + 1
1675 4 : rows = 0
1676 : END IF
1677 24 : rows = rows + atom_rows
1678 : END DO
1679 10 : IF (rows > 0) nsubchunks = nsubchunks + 1
1680 10 : nsubchunks = MAX(1, nsubchunks)
1681 :
1682 10 : END FUNCTION skala_gpw_atom_subchunk_count
1683 :
1684 : ! **************************************************************************************************
1685 : !> \brief Build an atom-contiguous subchunk feature bundle from a rank-local atom chunk.
1686 : !> \param parent ...
1687 : !> \param features ...
1688 : !> \param subchunk_index ...
1689 : !> \param max_rows ...
1690 : !> \param requires_grad ...
1691 : ! **************************************************************************************************
1692 4 : SUBROUTINE skala_gpw_feature_build_atom_subchunk(parent, features, subchunk_index, &
1693 : max_rows, requires_grad)
1694 : TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1695 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1696 : INTEGER, INTENT(IN) :: subchunk_index, max_rows
1697 : LOGICAL, INTENT(IN) :: requires_grad
1698 :
1699 : INTEGER :: atom_begin, atom_count, atom_end, &
1700 : max_grid_size, row_begin, row_count, &
1701 : row_end
1702 :
1703 4 : CALL skala_gpw_feature_release(features)
1704 4 : CPASSERT(parent%uses_atom_chunks)
1705 : CALL atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
1706 4 : row_begin, row_end)
1707 4 : atom_count = atom_end - atom_begin + 1
1708 4 : row_count = row_end - row_begin + 1
1709 4 : CPASSERT(atom_count > 0)
1710 4 : CPASSERT(row_count > 0)
1711 : MARK_USED(requires_grad)
1712 8 : max_grid_size = MAXVAL(INT(cached_layout%chunk_atomic_grid_sizes(atom_begin:atom_end)))
1713 :
1714 8 : ALLOCATE (features%atomic_grid_size_bound_shape(0, max_grid_size))
1715 64004 : features%atomic_grid_size_bound_shape = 0_int_8
1716 :
1717 4 : features%chunk_feature_count = row_count
1718 4 : features%nflat = parent%nflat
1719 4 : features%nflat_local = parent%nflat_local
1720 64004 : features%grid_weight_sum = SUM(cached_layout%chunk_grid_weights(row_begin:row_end))
1721 4 : features%uses_atom_chunks = .TRUE.
1722 4 : features%uses_atom_chunk_routing = parent%uses_atom_chunk_routing
1723 : CALL add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
1724 4 : row_count)
1725 4 : features%active = .TRUE.
1726 :
1727 4 : END SUBROUTINE skala_gpw_feature_build_atom_subchunk
1728 :
1729 : ! **************************************************************************************************
1730 : !> \brief Return atom and row bounds for an atom-contiguous rank-local subchunk.
1731 : !> \param subchunk_index ...
1732 : !> \param max_rows ...
1733 : !> \param atom_begin ...
1734 : !> \param atom_end ...
1735 : !> \param row_begin ...
1736 : !> \param row_end ...
1737 : ! **************************************************************************************************
1738 4 : SUBROUTINE atom_subchunk_bounds(subchunk_index, max_rows, atom_begin, atom_end, &
1739 : row_begin, row_end)
1740 : INTEGER, INTENT(IN) :: subchunk_index, max_rows
1741 : INTEGER, INTENT(OUT) :: atom_begin, atom_end, row_begin, row_end
1742 :
1743 : INTEGER :: atom_rows, current_subchunk, iatom, &
1744 : row_cursor, rows
1745 :
1746 4 : CPASSERT(subchunk_index > 0)
1747 4 : CPASSERT(max_rows > 0)
1748 4 : CPASSERT(cached_layout%chunk_natom > 0)
1749 :
1750 4 : atom_begin = 1
1751 4 : atom_end = 0
1752 4 : row_begin = 1
1753 4 : row_end = 0
1754 4 : current_subchunk = 1
1755 4 : row_cursor = 1
1756 4 : rows = 0
1757 10 : DO iatom = 1, cached_layout%chunk_natom
1758 8 : atom_rows = INT(cached_layout%chunk_atomic_grid_sizes(iatom))
1759 8 : IF (rows > 0 .AND. rows + atom_rows > max_rows) THEN
1760 4 : IF (current_subchunk == subchunk_index) THEN
1761 2 : atom_end = iatom - 1
1762 2 : row_end = row_cursor - 1
1763 2 : RETURN
1764 : END IF
1765 2 : current_subchunk = current_subchunk + 1
1766 2 : atom_begin = iatom
1767 2 : row_begin = row_cursor
1768 2 : rows = 0
1769 : END IF
1770 6 : rows = rows + atom_rows
1771 8 : row_cursor = row_cursor + atom_rows
1772 : END DO
1773 :
1774 2 : IF (current_subchunk == subchunk_index) THEN
1775 2 : atom_end = cached_layout%chunk_natom
1776 2 : row_end = row_cursor - 1
1777 2 : RETURN
1778 : END IF
1779 :
1780 0 : CPABORT("Requested native SKALA atom subchunk does not exist.")
1781 :
1782 : END SUBROUTINE atom_subchunk_bounds
1783 :
1784 : ! **************************************************************************************************
1785 : !> \brief Insert a subchunk into a Torch dictionary using static views of the cached chunk tensors.
1786 : !> \param parent ...
1787 : !> \param features ...
1788 : !> \param atom_begin ...
1789 : !> \param atom_count ...
1790 : !> \param row_begin ...
1791 : !> \param row_count ...
1792 : ! **************************************************************************************************
1793 4 : SUBROUTINE add_subchunk_feature_tensors(parent, features, atom_begin, atom_count, row_begin, &
1794 : row_count)
1795 : TYPE(skala_gpw_feature_type), INTENT(IN) :: parent
1796 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1797 : INTEGER, INTENT(IN) :: atom_begin, atom_count, row_begin, &
1798 : row_count
1799 :
1800 4 : CPASSERT(cached_layout%chunk_static_tensors_active)
1801 4 : CPASSERT(parent%active)
1802 4 : CPASSERT(ALLOCATED(features%atomic_grid_size_bound_shape))
1803 :
1804 4 : features%owns_coordinate_tensor = .FALSE.
1805 4 : features%owns_dynamic_tensors = .TRUE.
1806 4 : features%owns_inputs = .TRUE.
1807 4 : features%owns_static_tensors = .TRUE.
1808 4 : features%uses_collapsed_rks_dynamic = parent%uses_collapsed_rks_dynamic
1809 :
1810 : CALL torch_tensor_narrow(cached_layout%chunk_grid_coords_t, 0, row_begin - 1, &
1811 4 : row_count, features%grid_coords_t)
1812 : CALL torch_tensor_narrow(cached_layout%chunk_grid_weights_t, 0, row_begin - 1, &
1813 4 : row_count, features%grid_weights_t)
1814 : CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_weights_t, 0, row_begin - 1, &
1815 4 : row_count, features%atomic_grid_weights_t)
1816 : CALL torch_tensor_narrow(cached_layout%chunk_atomic_grid_sizes_t, 0, atom_begin - 1, &
1817 4 : atom_count, features%atomic_grid_sizes_t)
1818 : CALL torch_tensor_narrow(cached_layout%chunk_coarse_0_atomic_coords_t, 0, &
1819 4 : atom_begin - 1, atom_count, features%coarse_0_atomic_coords_t)
1820 : CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
1821 4 : features%atomic_grid_size_bound_shape)
1822 4 : CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .FALSE.)
1823 : CALL torch_tensor_narrow(parent%density_t, 1, row_begin - 1, row_count, &
1824 4 : features%density_t)
1825 4 : CALL torch_tensor_narrow(parent%grad_t, 2, row_begin - 1, row_count, features%grad_t)
1826 4 : CALL torch_tensor_narrow(parent%kin_t, 1, row_begin - 1, row_count, features%kin_t)
1827 4 : IF (features%uses_collapsed_rks_dynamic) THEN
1828 4 : CALL torch_tensor_expand_dim(features%density_t, 0, 2, features%density_input_t)
1829 4 : CALL torch_tensor_expand_dim(features%grad_t, 0, 2, features%grad_input_t)
1830 4 : CALL torch_tensor_expand_dim(features%kin_t, 0, 2, features%kin_input_t)
1831 : END IF
1832 :
1833 4 : CALL torch_dict_create(features%inputs)
1834 4 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
1835 4 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
1836 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
1837 4 : features%atomic_grid_weights_t)
1838 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
1839 4 : features%atomic_grid_sizes_t)
1840 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
1841 4 : features%atomic_grid_size_bound_shape_t)
1842 4 : IF (features%uses_collapsed_rks_dynamic) THEN
1843 4 : CALL torch_dict_insert(features%inputs, "density", features%density_input_t)
1844 4 : CALL torch_dict_insert(features%inputs, "grad", features%grad_input_t)
1845 4 : CALL torch_dict_insert(features%inputs, "kin", features%kin_input_t)
1846 : ELSE
1847 0 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
1848 0 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
1849 0 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
1850 : END IF
1851 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
1852 4 : features%coarse_0_atomic_coords_t)
1853 :
1854 4 : END SUBROUTINE add_subchunk_feature_tensors
1855 :
1856 : ! **************************************************************************************************
1857 : !> \brief Insert owned subchunk arrays into a Torch dictionary.
1858 : !> \param features ...
1859 : !> \param requires_grad ...
1860 : ! **************************************************************************************************
1861 0 : SUBROUTINE add_owned_feature_tensors(features, requires_grad)
1862 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1863 : LOGICAL, INTENT(IN) :: requires_grad
1864 :
1865 0 : CPASSERT(ALLOCATED(features%chunk_density))
1866 0 : CPASSERT(ALLOCATED(features%chunk_grad))
1867 0 : CPASSERT(ALLOCATED(features%chunk_kin))
1868 0 : CPASSERT(ALLOCATED(features%grid_coords))
1869 0 : CPASSERT(ALLOCATED(features%grid_weights))
1870 0 : CPASSERT(ALLOCATED(features%atomic_grid_weights))
1871 0 : CPASSERT(ALLOCATED(features%atomic_grid_sizes))
1872 0 : CPASSERT(ALLOCATED(features%atomic_grid_size_bound_shape))
1873 0 : CPASSERT(ALLOCATED(features%coarse_0_atomic_coords))
1874 :
1875 0 : features%owns_coordinate_tensor = .FALSE.
1876 0 : features%owns_dynamic_tensors = .TRUE.
1877 0 : features%owns_inputs = .TRUE.
1878 0 : features%owns_static_tensors = .TRUE.
1879 :
1880 0 : CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
1881 0 : CALL torch_tensor_to_device_leaf(features%grid_coords_t, .FALSE.)
1882 0 : CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
1883 0 : CALL torch_tensor_to_device_leaf(features%grid_weights_t, .FALSE.)
1884 0 : CALL torch_tensor_from_array(features%atomic_grid_weights_t, features%atomic_grid_weights)
1885 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .FALSE.)
1886 0 : CALL torch_tensor_from_array(features%atomic_grid_sizes_t, features%atomic_grid_sizes)
1887 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_sizes_t, .FALSE.)
1888 : CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
1889 0 : features%coarse_0_atomic_coords)
1890 0 : CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .FALSE.)
1891 : CALL torch_tensor_from_array(features%atomic_grid_size_bound_shape_t, &
1892 0 : features%atomic_grid_size_bound_shape)
1893 0 : CALL torch_tensor_to_device_leaf(features%atomic_grid_size_bound_shape_t, .FALSE.)
1894 0 : CALL torch_tensor_from_array(features%density_t, features%chunk_density)
1895 0 : CALL torch_tensor_to_device_leaf(features%density_t, requires_grad)
1896 0 : CALL torch_tensor_from_array(features%grad_t, features%chunk_grad)
1897 0 : CALL torch_tensor_to_device_leaf(features%grad_t, requires_grad)
1898 0 : CALL torch_tensor_from_array(features%kin_t, features%chunk_kin)
1899 0 : CALL torch_tensor_to_device_leaf(features%kin_t, requires_grad)
1900 :
1901 0 : CALL torch_dict_create(features%inputs)
1902 0 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
1903 0 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
1904 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
1905 0 : features%atomic_grid_weights_t)
1906 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
1907 0 : features%atomic_grid_sizes_t)
1908 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
1909 0 : features%atomic_grid_size_bound_shape_t)
1910 0 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
1911 0 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
1912 0 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
1913 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
1914 0 : features%coarse_0_atomic_coords_t)
1915 :
1916 0 : END SUBROUTINE add_owned_feature_tensors
1917 :
1918 : ! **************************************************************************************************
1919 : !> \brief Insert all SKALA feature tensors into the Torch dictionary.
1920 : !> \param features ...
1921 : !> \param requires_grad ...
1922 : !> \param requires_coordinate_grad ...
1923 : !> \param requires_stress_grad ...
1924 : !> \param use_atom_chunks ...
1925 : !> \param requires_weight_grad ...
1926 : ! **************************************************************************************************
1927 290 : SUBROUTINE add_feature_tensors(features, requires_grad, requires_coordinate_grad, &
1928 : requires_stress_grad, use_atom_chunks, requires_weight_grad)
1929 : TYPE(skala_gpw_feature_type), INTENT(INOUT) :: features
1930 : LOGICAL, INTENT(IN) :: requires_grad, requires_coordinate_grad, &
1931 : requires_stress_grad, use_atom_chunks
1932 : LOGICAL, INTENT(IN), OPTIONAL :: requires_weight_grad
1933 :
1934 : LOGICAL :: my_requires_weight_grad
1935 :
1936 290 : my_requires_weight_grad = .FALSE.
1937 290 : IF (PRESENT(requires_weight_grad)) my_requires_weight_grad = requires_weight_grad
1938 :
1939 290 : CPASSERT(cached_layout%static_tensors_active)
1940 290 : features%owns_static_tensors = .FALSE.
1941 290 : features%owns_coordinate_tensor = .FALSE.
1942 290 : features%owns_grid_coordinate_tensor = .FALSE.
1943 290 : features%owns_weight_tensors = .FALSE.
1944 290 : features%owns_dynamic_tensors = .FALSE.
1945 290 : features%owns_inputs = .TRUE.
1946 290 : IF (use_atom_chunks) THEN
1947 8 : CPASSERT(.NOT. requires_coordinate_grad)
1948 8 : CPASSERT(.NOT. requires_stress_grad)
1949 8 : CPASSERT(.NOT. my_requires_weight_grad)
1950 8 : CPASSERT(cached_layout%chunk_static_tensors_active)
1951 8 : features%grid_coords_t = cached_layout%chunk_grid_coords_t
1952 8 : features%grid_weights_t = cached_layout%chunk_grid_weights_t
1953 8 : features%atomic_grid_weights_t = cached_layout%chunk_atomic_grid_weights_t
1954 8 : features%atomic_grid_sizes_t = cached_layout%chunk_atomic_grid_sizes_t
1955 : features%atomic_grid_size_bound_shape_t = &
1956 8 : cached_layout%chunk_atomic_grid_size_bound_shape_t
1957 8 : features%local_feature_indices_t = cached_layout%chunk_feature_indices_t
1958 :
1959 8 : IF (cached_layout%chunk_inputs_active .AND. &
1960 : (cached_layout%chunk_inputs_use_collapsed_rks .NEQV. &
1961 : features%uses_collapsed_rks_dynamic)) THEN
1962 0 : CALL torch_dict_release(cached_layout%chunk_inputs)
1963 0 : cached_layout%chunk_inputs_active = .FALSE.
1964 : END IF
1965 8 : IF (.NOT. features%uses_collapsed_rks_dynamic .AND. &
1966 : cached_layout%chunk_dynamic_input_views_active) THEN
1967 0 : CALL torch_tensor_release(cached_layout%chunk_density_input_t)
1968 0 : CALL torch_tensor_release(cached_layout%chunk_grad_input_t)
1969 0 : CALL torch_tensor_release(cached_layout%chunk_kin_input_t)
1970 0 : cached_layout%chunk_dynamic_input_views_active = .FALSE.
1971 : END IF
1972 :
1973 : CALL torch_tensor_reset_from_array(cached_layout%chunk_density_t, &
1974 8 : features%chunk_density, requires_grad=requires_grad)
1975 8 : features%density_t = cached_layout%chunk_density_t
1976 : CALL torch_tensor_reset_from_array(cached_layout%chunk_grad_t, features%chunk_grad, &
1977 8 : requires_grad=requires_grad)
1978 8 : features%grad_t = cached_layout%chunk_grad_t
1979 : CALL torch_tensor_reset_from_array(cached_layout%chunk_kin_t, features%chunk_kin, &
1980 8 : requires_grad=requires_grad)
1981 8 : features%kin_t = cached_layout%chunk_kin_t
1982 8 : cached_layout%chunk_dynamic_tensors_active = .TRUE.
1983 :
1984 8 : IF (features%uses_collapsed_rks_dynamic .AND. &
1985 : .NOT. cached_layout%chunk_dynamic_input_views_active) THEN
1986 : CALL torch_tensor_expand_dim(cached_layout%chunk_density_t, 0, 2, &
1987 6 : cached_layout%chunk_density_input_t)
1988 : CALL torch_tensor_expand_dim(cached_layout%chunk_grad_t, 0, 2, &
1989 6 : cached_layout%chunk_grad_input_t)
1990 : CALL torch_tensor_expand_dim(cached_layout%chunk_kin_t, 0, 2, &
1991 6 : cached_layout%chunk_kin_input_t)
1992 6 : cached_layout%chunk_dynamic_input_views_active = .TRUE.
1993 : END IF
1994 8 : IF (features%uses_collapsed_rks_dynamic) THEN
1995 8 : features%density_input_t = cached_layout%chunk_density_input_t
1996 8 : features%grad_input_t = cached_layout%chunk_grad_input_t
1997 8 : features%kin_input_t = cached_layout%chunk_kin_input_t
1998 : END IF
1999 :
2000 8 : IF (.NOT. cached_layout%chunk_inputs_active) THEN
2001 6 : CALL torch_dict_clone(cached_layout%chunk_static_inputs, cached_layout%chunk_inputs)
2002 6 : IF (features%uses_collapsed_rks_dynamic) THEN
2003 : CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2004 6 : features%density_input_t)
2005 : CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2006 6 : features%grad_input_t)
2007 : CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2008 6 : features%kin_input_t)
2009 : ELSE
2010 : CALL torch_dict_insert(cached_layout%chunk_inputs, "density", &
2011 0 : cached_layout%chunk_density_t)
2012 : CALL torch_dict_insert(cached_layout%chunk_inputs, "grad", &
2013 0 : cached_layout%chunk_grad_t)
2014 : CALL torch_dict_insert(cached_layout%chunk_inputs, "kin", &
2015 0 : cached_layout%chunk_kin_t)
2016 : END IF
2017 : CALL torch_dict_insert(cached_layout%chunk_inputs, "coarse_0_atomic_coords", &
2018 6 : cached_layout%chunk_coarse_0_atomic_coords_t)
2019 6 : cached_layout%chunk_inputs_use_collapsed_rks = features%uses_collapsed_rks_dynamic
2020 6 : cached_layout%chunk_inputs_active = .TRUE.
2021 : END IF
2022 8 : features%inputs = cached_layout%chunk_inputs
2023 8 : features%owns_inputs = .FALSE.
2024 8 : features%coarse_0_atomic_coords_t = cached_layout%chunk_coarse_0_atomic_coords_t
2025 : ELSE
2026 282 : IF (.NOT. requires_stress_grad .AND. .NOT. my_requires_weight_grad) THEN
2027 222 : features%grid_coords_t = cached_layout%grid_coords_t
2028 222 : features%grid_weights_t = cached_layout%grid_weights_t
2029 222 : features%atomic_grid_weights_t = cached_layout%atomic_grid_weights_t
2030 : END IF
2031 282 : features%atomic_grid_sizes_t = cached_layout%atomic_grid_sizes_t
2032 282 : features%atomic_grid_size_bound_shape_t = cached_layout%atomic_grid_size_bound_shape_t
2033 282 : features%local_feature_indices_t = cached_layout%local_feature_indices_t
2034 :
2035 : CALL torch_tensor_reset_from_array(cached_layout%density_t, features%density, &
2036 282 : requires_grad=requires_grad)
2037 282 : features%density_t = cached_layout%density_t
2038 : CALL torch_tensor_reset_from_array(cached_layout%grad_t, features%grad, &
2039 282 : requires_grad=requires_grad)
2040 282 : features%grad_t = cached_layout%grad_t
2041 : CALL torch_tensor_reset_from_array(cached_layout%kin_t, features%kin, &
2042 282 : requires_grad=requires_grad)
2043 282 : features%kin_t = cached_layout%kin_t
2044 282 : cached_layout%dynamic_tensors_active = .TRUE.
2045 :
2046 282 : IF (requires_coordinate_grad .OR. requires_stress_grad .OR. my_requires_weight_grad) THEN
2047 60 : IF (requires_stress_grad .OR. my_requires_weight_grad) THEN
2048 60 : CALL torch_dict_create(features%inputs)
2049 60 : IF (requires_stress_grad) THEN
2050 50 : CALL torch_tensor_from_array(features%grid_coords_t, features%grid_coords)
2051 50 : CALL torch_tensor_to_device_leaf(features%grid_coords_t, .TRUE.)
2052 50 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2053 50 : features%owns_grid_coordinate_tensor = .TRUE.
2054 : ELSE
2055 10 : features%grid_coords_t = cached_layout%grid_coords_t
2056 10 : CALL torch_dict_insert(features%inputs, "grid_coords", features%grid_coords_t)
2057 : END IF
2058 60 : CALL torch_tensor_from_array(features%grid_weights_t, features%grid_weights)
2059 60 : CALL torch_tensor_to_device_leaf(features%grid_weights_t, .TRUE.)
2060 : CALL torch_tensor_from_array(features%atomic_grid_weights_t, &
2061 60 : features%atomic_grid_weights)
2062 60 : CALL torch_tensor_to_device_leaf(features%atomic_grid_weights_t, .TRUE.)
2063 60 : CALL torch_dict_insert(features%inputs, "grid_weights", features%grid_weights_t)
2064 : CALL torch_dict_insert(features%inputs, "atomic_grid_weights", &
2065 60 : features%atomic_grid_weights_t)
2066 : CALL torch_dict_insert(features%inputs, "atomic_grid_sizes", &
2067 60 : features%atomic_grid_sizes_t)
2068 : CALL torch_dict_insert(features%inputs, "atomic_grid_size_bound_shape", &
2069 60 : features%atomic_grid_size_bound_shape_t)
2070 60 : features%owns_weight_tensors = .TRUE.
2071 : ELSE
2072 0 : CALL torch_dict_clone(cached_layout%static_inputs, features%inputs)
2073 : END IF
2074 60 : CALL torch_dict_insert(features%inputs, "density", features%density_t)
2075 60 : CALL torch_dict_insert(features%inputs, "grad", features%grad_t)
2076 60 : CALL torch_dict_insert(features%inputs, "kin", features%kin_t)
2077 : ELSE
2078 222 : IF (.NOT. cached_layout%inputs_active) THEN
2079 122 : CALL torch_dict_clone(cached_layout%static_inputs, cached_layout%inputs)
2080 122 : CALL torch_dict_insert(cached_layout%inputs, "density", cached_layout%density_t)
2081 122 : CALL torch_dict_insert(cached_layout%inputs, "grad", cached_layout%grad_t)
2082 122 : CALL torch_dict_insert(cached_layout%inputs, "kin", cached_layout%kin_t)
2083 : CALL torch_dict_insert(cached_layout%inputs, "coarse_0_atomic_coords", &
2084 122 : cached_layout%coarse_0_atomic_coords_t)
2085 122 : cached_layout%inputs_active = .TRUE.
2086 : END IF
2087 222 : features%inputs = cached_layout%inputs
2088 222 : features%owns_inputs = .FALSE.
2089 222 : features%coarse_0_atomic_coords_t = cached_layout%coarse_0_atomic_coords_t
2090 : END IF
2091 : END IF
2092 :
2093 290 : IF (requires_coordinate_grad .OR. requires_stress_grad) THEN
2094 60 : CPASSERT(.NOT. use_atom_chunks)
2095 : CALL torch_tensor_from_array(features%coarse_0_atomic_coords_t, &
2096 60 : features%coarse_0_atomic_coords)
2097 60 : CALL torch_tensor_to_device_leaf(features%coarse_0_atomic_coords_t, .TRUE.)
2098 : CALL torch_dict_insert(features%inputs, "coarse_0_atomic_coords", &
2099 60 : features%coarse_0_atomic_coords_t)
2100 60 : features%owns_coordinate_tensor = .TRUE.
2101 : END IF
2102 :
2103 290 : END SUBROUTINE add_feature_tensors
2104 :
2105 : ! **************************************************************************************************
2106 : !> \brief Return the Cartesian coordinate of a regular GPW grid point.
2107 : !> \param pw_grid ...
2108 : !> \param index ...
2109 : !> \return ...
2110 : ! **************************************************************************************************
2111 1364695 : FUNCTION grid_coordinate(pw_grid, index) RESULT(coord)
2112 : TYPE(pw_grid_type), POINTER :: pw_grid
2113 : INTEGER, DIMENSION(3), INTENT(IN) :: index
2114 : REAL(KIND=dp), DIMENSION(3) :: coord
2115 :
2116 : INTEGER, DIMENSION(3) :: relative_index
2117 :
2118 5458780 : relative_index = index - pw_grid%bounds(1, :)
2119 : coord = REAL(relative_index(1), KIND=dp)*pw_grid%dh(:, 1) + &
2120 : REAL(relative_index(2), KIND=dp)*pw_grid%dh(:, 2) + &
2121 5458780 : REAL(relative_index(3), KIND=dp)*pw_grid%dh(:, 3)
2122 :
2123 1364695 : END FUNCTION grid_coordinate
2124 :
2125 : ! **************************************************************************************************
2126 : !> \brief Build Becke-like smooth atom weights for one native-grid point.
2127 : !> \param grid_point ...
2128 : !> \param atom_coords ...
2129 : !> \param cell ...
2130 : !> \param weights ...
2131 : !> \param atom_image_coords ...
2132 : !> \param distances ...
2133 : ! **************************************************************************************************
2134 377508 : SUBROUTINE smooth_atom_partition(grid_point, atom_coords, cell, weights, atom_image_coords, &
2135 377508 : distances)
2136 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2137 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2138 : TYPE(cell_type), POINTER :: cell
2139 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: weights
2140 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT) :: atom_image_coords
2141 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: distances
2142 :
2143 : INTEGER :: iatom, jatom, natom
2144 : REAL(KIND=dp) :: mu, rab, rsum, switch, total
2145 : REAL(KIND=dp), DIMENSION(3) :: rij
2146 755016 : REAL(KIND=dp), DIMENSION(3, SIZE(atom_coords, 2)) :: partition_atom_coords
2147 :
2148 377508 : natom = SIZE(atom_coords, 2)
2149 377508 : CPASSERT(SIZE(weights) == natom)
2150 377508 : CPASSERT(SIZE(atom_image_coords, 1) == 3)
2151 377508 : CPASSERT(SIZE(atom_image_coords, 2) == natom)
2152 377508 : CPASSERT(SIZE(distances) == natom)
2153 :
2154 1132524 : DO iatom = 1, natom
2155 : atom_image_coords(:, iatom) = &
2156 755016 : nearest_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2157 : partition_atom_coords(:, iatom) = &
2158 755016 : nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2159 3020064 : rij = grid_point - partition_atom_coords(:, iatom)
2160 3397572 : distances(iatom) = SQRT(SUM(rij**2))
2161 : END DO
2162 :
2163 1132524 : weights = 1.0_dp
2164 755016 : DO iatom = 1, natom - 1
2165 1132524 : DO jatom = iatom + 1, natom
2166 1510032 : rij = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
2167 1510032 : rab = SQRT(SUM(rij**2))
2168 377508 : IF (rab <= layout_tol) CYCLE
2169 377508 : mu = (distances(iatom) - distances(jatom))/rab
2170 377508 : mu = MAX(-1.0_dp, MIN(1.0_dp, mu))
2171 377508 : switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2172 377508 : weights(iatom) = weights(iatom)*switch
2173 755016 : weights(jatom) = weights(jatom)*(1.0_dp - switch)
2174 : END DO
2175 : END DO
2176 :
2177 1132524 : total = SUM(weights)
2178 377508 : IF (total > 0.0_dp) THEN
2179 1132524 : weights = weights/total
2180 : ELSE
2181 : rsum = HUGE(1.0_dp)
2182 : jatom = 1
2183 0 : DO iatom = 1, natom
2184 0 : IF (distances(iatom) < rsum) THEN
2185 0 : rsum = distances(iatom)
2186 0 : jatom = iatom
2187 : END IF
2188 : END DO
2189 0 : weights = 0.0_dp
2190 0 : weights(jatom) = 1.0_dp
2191 : END IF
2192 :
2193 377508 : END SUBROUTINE smooth_atom_partition
2194 :
2195 : ! **************************************************************************************************
2196 : !> \brief Build smooth atom weights and their atom/cell deformation derivatives.
2197 : !> \param grid_point ...
2198 : !> \param atom_coords ...
2199 : !> \param cell ...
2200 : !> \param weights ...
2201 : !> \param included ...
2202 : !> \param dweights_datom ...
2203 : !> \param dweights_dstrain ...
2204 : ! **************************************************************************************************
2205 554595 : SUBROUTINE skala_gpw_smooth_partition_derivatives(grid_point, atom_coords, cell, &
2206 554595 : weights, included, dweights_datom, &
2207 554595 : dweights_dstrain)
2208 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2209 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2210 : TYPE(cell_type), POINTER :: cell
2211 : REAL(KIND=dp), DIMENSION(:), INTENT(OUT) :: weights
2212 : LOGICAL, DIMENSION(:), INTENT(OUT) :: included
2213 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(OUT) :: dweights_datom, dweights_dstrain
2214 :
2215 : INTEGER :: iatom, idir, jatom, jdir, natom
2216 : REAL(KIND=dp) :: dist_diff, ds_dmu, included_sum, mu, &
2217 : mu_raw, one_minus_switch, rab, rsum, &
2218 : switch, total
2219 : REAL(KIND=dp), DIMENSION(3) :: dmu_atom_i, dmu_atom_j, ds_atom_i, &
2220 : ds_atom_j, pair, unit_pair
2221 : REAL(KIND=dp), DIMENSION(3, 3) :: dmu_strain, ds_strain, mean_strain
2222 : REAL(KIND=dp), DIMENSION(3, SIZE(atom_coords, 2), &
2223 1109190 : SIZE(atom_coords, 2)) :: log_weight_atom
2224 1109190 : REAL(KIND=dp), DIMENSION(3, SIZE(atom_coords, 2)) :: mean_atom, partition_atom_coords, rvecs, &
2225 1109190 : unit_rvecs
2226 : REAL(KIND=dp), &
2227 1109190 : DIMENSION(3, 3, SIZE(atom_coords, 2)) :: log_weight_strain
2228 1109190 : REAL(KIND=dp), DIMENSION(SIZE(atom_coords, 2)) :: distances, normalized_weights, &
2229 554595 : raw_weights
2230 :
2231 554595 : natom = SIZE(atom_coords, 2)
2232 554595 : CPASSERT(SIZE(weights) == natom)
2233 554595 : CPASSERT(SIZE(included) == natom)
2234 554595 : CPASSERT(SIZE(dweights_datom, 1) == 3)
2235 554595 : CPASSERT(SIZE(dweights_datom, 2) == natom)
2236 554595 : CPASSERT(SIZE(dweights_datom, 3) == natom)
2237 554595 : CPASSERT(SIZE(dweights_dstrain, 1) == 3)
2238 554595 : CPASSERT(SIZE(dweights_dstrain, 2) == 3)
2239 554595 : CPASSERT(SIZE(dweights_dstrain, 3) == natom)
2240 :
2241 1663785 : weights = 0.0_dp
2242 1663785 : included = .FALSE.
2243 10537305 : dweights_datom = 0.0_dp
2244 14974065 : dweights_dstrain = 0.0_dp
2245 1663785 : raw_weights = 1.0_dp
2246 10537305 : log_weight_atom = 0.0_dp
2247 14974065 : log_weight_strain = 0.0_dp
2248 :
2249 1663785 : DO iatom = 1, natom
2250 : partition_atom_coords(:, iatom) = &
2251 1109190 : nearest_atom_image_coordinate(atom_coords(:, iatom), grid_point, cell)
2252 4436760 : rvecs(:, iatom) = grid_point - partition_atom_coords(:, iatom)
2253 4436760 : distances(iatom) = SQRT(SUM(rvecs(:, iatom)**2))
2254 1663785 : IF (distances(iatom) > layout_tol) THEN
2255 4436760 : unit_rvecs(:, iatom) = rvecs(:, iatom)/distances(iatom)
2256 : ELSE
2257 0 : unit_rvecs(:, iatom) = 0.0_dp
2258 : END IF
2259 : END DO
2260 :
2261 1109190 : DO iatom = 1, natom - 1
2262 1663785 : DO jatom = iatom + 1, natom
2263 2218380 : pair = partition_atom_coords(:, iatom) - partition_atom_coords(:, jatom)
2264 2218380 : rab = SQRT(SUM(pair**2))
2265 554595 : IF (rab <= layout_tol) CYCLE
2266 2218380 : unit_pair = pair/rab
2267 554595 : dist_diff = distances(iatom) - distances(jatom)
2268 554595 : mu_raw = dist_diff/rab
2269 554595 : mu = MAX(-1.0_dp, MIN(1.0_dp, mu_raw))
2270 554595 : switch = 0.5_dp*(1.0_dp - becke_shape(mu))
2271 554595 : one_minus_switch = 1.0_dp - switch
2272 :
2273 554595 : IF (ABS(mu_raw) < 1.0_dp) THEN
2274 554133 : ds_dmu = -0.5_dp*becke_shape_derivative(mu)
2275 : ELSE
2276 : ds_dmu = 0.0_dp
2277 : END IF
2278 554133 : IF (ABS(ds_dmu) > 0.0_dp .AND. switch > TINY(1.0_dp) .AND. &
2279 : one_minus_switch > TINY(1.0_dp)) THEN
2280 2215644 : dmu_atom_i = (-unit_rvecs(:, iatom)*rab - dist_diff*unit_pair)/rab**2
2281 2215644 : dmu_atom_j = (unit_rvecs(:, jatom)*rab + dist_diff*unit_pair)/rab**2
2282 2215644 : ds_atom_i = ds_dmu*dmu_atom_i
2283 2215644 : ds_atom_j = ds_dmu*dmu_atom_j
2284 : log_weight_atom(:, iatom, iatom) = &
2285 2215644 : log_weight_atom(:, iatom, iatom) + ds_atom_i/switch
2286 : log_weight_atom(:, iatom, jatom) = &
2287 2215644 : log_weight_atom(:, iatom, jatom) - ds_atom_i/one_minus_switch
2288 : log_weight_atom(:, jatom, iatom) = &
2289 2215644 : log_weight_atom(:, jatom, iatom) + ds_atom_j/switch
2290 : log_weight_atom(:, jatom, jatom) = &
2291 2215644 : log_weight_atom(:, jatom, jatom) - ds_atom_j/one_minus_switch
2292 :
2293 2215644 : DO idir = 1, 3
2294 7200843 : DO jdir = 1, 3
2295 : dmu_strain(idir, jdir) = &
2296 : ((unit_rvecs(idir, iatom)*rvecs(jdir, iatom) - &
2297 : unit_rvecs(idir, jatom)*rvecs(jdir, jatom))*rab - &
2298 6646932 : dist_diff*unit_pair(idir)*pair(jdir))/rab**2
2299 : END DO
2300 : END DO
2301 7200843 : ds_strain = ds_dmu*dmu_strain
2302 : log_weight_strain(:, :, iatom) = &
2303 7200843 : log_weight_strain(:, :, iatom) + ds_strain/switch
2304 : log_weight_strain(:, :, jatom) = &
2305 7200843 : log_weight_strain(:, :, jatom) - ds_strain/one_minus_switch
2306 : END IF
2307 :
2308 554595 : raw_weights(iatom) = raw_weights(iatom)*switch
2309 1109190 : raw_weights(jatom) = raw_weights(jatom)*one_minus_switch
2310 : END DO
2311 : END DO
2312 :
2313 1663785 : total = SUM(raw_weights)
2314 554595 : IF (total > 0.0_dp) THEN
2315 1663785 : normalized_weights = raw_weights/total
2316 1663785 : included = normalized_weights > smooth_partition_eps
2317 : ELSE
2318 : rsum = HUGE(1.0_dp)
2319 : jatom = 1
2320 0 : DO iatom = 1, natom
2321 0 : IF (distances(iatom) < rsum) THEN
2322 0 : rsum = distances(iatom)
2323 0 : jatom = iatom
2324 : END IF
2325 : END DO
2326 0 : included(jatom) = .TRUE.
2327 0 : weights(jatom) = 1.0_dp
2328 0 : RETURN
2329 : END IF
2330 :
2331 1663785 : included_sum = SUM(raw_weights, MASK=included)
2332 554595 : IF (included_sum <= 0.0_dp) THEN
2333 : rsum = HUGE(1.0_dp)
2334 : jatom = 1
2335 0 : DO iatom = 1, natom
2336 0 : IF (distances(iatom) < rsum) THEN
2337 0 : rsum = distances(iatom)
2338 0 : jatom = iatom
2339 : END IF
2340 : END DO
2341 0 : included = .FALSE.
2342 0 : included(jatom) = .TRUE.
2343 0 : weights = 0.0_dp
2344 0 : weights(jatom) = 1.0_dp
2345 0 : RETURN
2346 : END IF
2347 :
2348 1663785 : DO iatom = 1, natom
2349 1663785 : IF (included(iatom)) weights(iatom) = raw_weights(iatom)/included_sum
2350 : END DO
2351 :
2352 4991355 : mean_atom = 0.0_dp
2353 554595 : mean_strain = 0.0_dp
2354 1663785 : DO iatom = 1, natom
2355 1109190 : IF (.NOT. included(iatom)) CYCLE
2356 14385774 : mean_strain = mean_strain + weights(iatom)*log_weight_strain(:, :, iatom)
2357 3874389 : DO jatom = 1, natom
2358 : mean_atom(:, jatom) = mean_atom(:, jatom) + &
2359 9961974 : weights(iatom)*log_weight_atom(:, jatom, iatom)
2360 : END DO
2361 : END DO
2362 :
2363 1663785 : DO iatom = 1, natom
2364 1109190 : IF (.NOT. included(iatom)) CYCLE
2365 : dweights_dstrain(:, :, iatom) = &
2366 14385774 : weights(iatom)*(log_weight_strain(:, :, iatom) - mean_strain)
2367 3874389 : DO jatom = 1, natom
2368 : dweights_datom(:, jatom, iatom) = &
2369 9961974 : weights(iatom)*(log_weight_atom(:, jatom, iatom) - mean_atom(:, jatom))
2370 : END DO
2371 : END DO
2372 :
2373 : END SUBROUTINE skala_gpw_smooth_partition_derivatives
2374 :
2375 : ! **************************************************************************************************
2376 : !> \brief Becke fuzzy-cell shape function.
2377 : !> \param mu ...
2378 : !> \return ...
2379 : ! **************************************************************************************************
2380 932103 : PURE FUNCTION becke_shape(mu) RESULT(val)
2381 : REAL(KIND=dp), INTENT(IN) :: mu
2382 : REAL(KIND=dp) :: val
2383 :
2384 : INTEGER :: iter
2385 :
2386 932103 : val = mu
2387 3728412 : DO iter = 1, 3
2388 3728412 : val = 0.5_dp*val*(3.0_dp - val*val)
2389 : END DO
2390 :
2391 932103 : END FUNCTION becke_shape
2392 :
2393 : ! **************************************************************************************************
2394 : !> \brief Derivative of the Becke fuzzy-cell shape function.
2395 : !> \param mu ...
2396 : !> \return ...
2397 : ! **************************************************************************************************
2398 554133 : PURE FUNCTION becke_shape_derivative(mu) RESULT(val)
2399 : REAL(KIND=dp), INTENT(IN) :: mu
2400 : REAL(KIND=dp) :: val
2401 :
2402 : INTEGER :: iter
2403 : REAL(KIND=dp) :: x
2404 :
2405 554133 : x = mu
2406 554133 : val = 1.0_dp
2407 2216532 : DO iter = 1, 3
2408 1662399 : val = val*1.5_dp*(1.0_dp - x*x)
2409 2216532 : x = 0.5_dp*x*(3.0_dp - x*x)
2410 : END DO
2411 :
2412 554133 : END FUNCTION becke_shape_derivative
2413 :
2414 : ! **************************************************************************************************
2415 : !> \brief Return the atom image nearest to a regular-grid point.
2416 : !> \param atom_coord ...
2417 : !> \param grid_point ...
2418 : !> \param cell ...
2419 : !> \return ...
2420 : ! **************************************************************************************************
2421 1864206 : FUNCTION nearest_atom_image_coordinate(atom_coord, grid_point, cell) RESULT(coord)
2422 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: atom_coord, grid_point
2423 : TYPE(cell_type), POINTER :: cell
2424 : REAL(KIND=dp), DIMENSION(3) :: coord
2425 :
2426 : REAL(KIND=dp) :: dx, dy, dz
2427 :
2428 1864206 : IF (cell%orthorhombic) THEN
2429 1864206 : dx = atom_coord(1) - grid_point(1)
2430 1864206 : dy = atom_coord(2) - grid_point(2)
2431 1864206 : dz = atom_coord(3) - grid_point(3)
2432 1864206 : dx = dx - cell%hmat(1, 1)*cell%perd(1)*ANINT(cell%h_inv(1, 1)*dx)
2433 1864206 : dy = dy - cell%hmat(2, 2)*cell%perd(2)*ANINT(cell%h_inv(2, 2)*dy)
2434 1864206 : dz = dz - cell%hmat(3, 3)*cell%perd(3)*ANINT(cell%h_inv(3, 3)*dz)
2435 7456824 : coord = grid_point + [dx, dy, dz]
2436 : ELSE
2437 0 : coord = grid_point + pbc(grid_point, atom_coord, cell)
2438 : END IF
2439 :
2440 1864206 : END FUNCTION nearest_atom_image_coordinate
2441 :
2442 : ! **************************************************************************************************
2443 : !> \brief Return the grid-point image nearest to the owning atom coordinate.
2444 : !> \param owner_coord ...
2445 : !> \param grid_point ...
2446 : !> \param cell ...
2447 : !> \return ...
2448 : ! **************************************************************************************************
2449 755016 : FUNCTION nearest_image_coordinate(owner_coord, grid_point, cell) RESULT(coord)
2450 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: owner_coord, grid_point
2451 : TYPE(cell_type), POINTER :: cell
2452 : REAL(KIND=dp), DIMENSION(3) :: coord
2453 :
2454 : REAL(KIND=dp) :: dx, dy, dz
2455 :
2456 755016 : IF (cell%orthorhombic) THEN
2457 755016 : dx = grid_point(1) - owner_coord(1)
2458 755016 : dy = grid_point(2) - owner_coord(2)
2459 755016 : dz = grid_point(3) - owner_coord(3)
2460 755016 : dx = dx - cell%hmat(1, 1)*cell%perd(1)*ANINT(cell%h_inv(1, 1)*dx)
2461 755016 : dy = dy - cell%hmat(2, 2)*cell%perd(2)*ANINT(cell%h_inv(2, 2)*dy)
2462 755016 : dz = dz - cell%hmat(3, 3)*cell%perd(3)*ANINT(cell%h_inv(3, 3)*dz)
2463 3020064 : coord = owner_coord + [dx, dy, dz]
2464 : ELSE
2465 0 : coord = owner_coord + pbc(owner_coord, grid_point, cell)
2466 : END IF
2467 :
2468 755016 : END FUNCTION nearest_image_coordinate
2469 :
2470 : ! **************************************************************************************************
2471 : !> \brief Assign a grid point to the nearest periodic atom.
2472 : !> \param grid_point ...
2473 : !> \param atom_coords ...
2474 : !> \param cell ...
2475 : !> \return ...
2476 : ! **************************************************************************************************
2477 987187 : FUNCTION nearest_atom(grid_point, atom_coords, cell) RESULT(owner)
2478 : REAL(KIND=dp), DIMENSION(3), INTENT(IN) :: grid_point
2479 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: atom_coords
2480 : TYPE(cell_type), POINTER :: cell
2481 : INTEGER :: owner
2482 :
2483 : INTEGER :: iatom
2484 : REAL(KIND=dp) :: best_r2, dx, dy, dz, r2
2485 : REAL(KIND=dp), DIMENSION(3) :: rij
2486 :
2487 987187 : owner = 1
2488 987187 : best_r2 = HUGE(1.0_dp)
2489 987187 : IF (cell%orthorhombic) THEN
2490 3886904 : DO iatom = 1, SIZE(atom_coords, 2)
2491 2899717 : dx = grid_point(1) - atom_coords(1, iatom)
2492 2899717 : dy = grid_point(2) - atom_coords(2, iatom)
2493 2899717 : dz = grid_point(3) - atom_coords(3, iatom)
2494 2899717 : dx = dx - cell%hmat(1, 1)*cell%perd(1)*ANINT(cell%h_inv(1, 1)*dx)
2495 2899717 : dy = dy - cell%hmat(2, 2)*cell%perd(2)*ANINT(cell%h_inv(2, 2)*dy)
2496 2899717 : dz = dz - cell%hmat(3, 3)*cell%perd(3)*ANINT(cell%h_inv(3, 3)*dz)
2497 2899717 : r2 = dx*dx + dy*dy + dz*dz
2498 3886904 : IF (r2 < best_r2) THEN
2499 1773819 : best_r2 = r2
2500 1773819 : owner = iatom
2501 : END IF
2502 : END DO
2503 : ELSE
2504 0 : DO iatom = 1, SIZE(atom_coords, 2)
2505 0 : rij = pbc(grid_point, atom_coords(:, iatom), cell)
2506 0 : r2 = SUM(rij**2)
2507 0 : IF (r2 < best_r2) THEN
2508 0 : best_r2 = r2
2509 0 : owner = iatom
2510 : END IF
2511 : END DO
2512 : END IF
2513 :
2514 987187 : END FUNCTION nearest_atom
2515 :
2516 0 : END MODULE skala_gpw_features
|
