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