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