Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2025 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief tensor index and mapping to DBM index
10 : !> \author Patrick Seewald
11 : ! **************************************************************************************************
12 : MODULE dbt_index
13 : USE dbt_allocate_wrap, ONLY: allocate_any
14 : USE kinds, ONLY: int_8
15 : #include "../base/base_uses.f90"
16 : #:include "dbt_macros.fypp"
17 :
18 : IMPLICIT NONE
19 : PRIVATE
20 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_index'
21 :
22 : PUBLIC :: &
23 : combine_tensor_index, &
24 : combine_pgrid_index, &
25 : create_nd_to_2d_mapping, &
26 : destroy_nd_to_2d_mapping, &
27 : get_2d_indices_tensor, &
28 : get_2d_indices_pgrid, &
29 : dbt_get_mapping_info, &
30 : get_nd_indices_tensor, &
31 : get_nd_indices_pgrid, &
32 : nd_to_2d_mapping, &
33 : ndims_mapping, &
34 : split_tensor_index, &
35 : split_pgrid_index, &
36 : ndims_mapping_row, &
37 : ndims_mapping_column, &
38 : dbt_inverse_order, &
39 : permute_index
40 :
41 : TYPE nd_to_2d_mapping
42 : INTEGER :: ndim_nd = -1
43 : INTEGER :: ndim1_2d = -1
44 : INTEGER :: ndim2_2d = -1
45 :
46 : INTEGER, DIMENSION(:), ALLOCATABLE :: dims_nd
47 : INTEGER(KIND=int_8), DIMENSION(2) :: dims_2d = -1
48 : INTEGER, DIMENSION(:), ALLOCATABLE :: dims1_2d
49 : INTEGER, DIMENSION(:), ALLOCATABLE :: dims2_2d
50 :
51 : INTEGER, DIMENSION(:), ALLOCATABLE :: map1_2d
52 : INTEGER, DIMENSION(:), ALLOCATABLE :: map2_2d
53 : INTEGER, DIMENSION(:), ALLOCATABLE :: map_nd
54 :
55 : INTEGER :: base = -1
56 : LOGICAL :: col_major = .FALSE.
57 : END TYPE nd_to_2d_mapping
58 :
59 : CONTAINS
60 :
61 : ! **************************************************************************************************
62 : !> \brief Create all data needed to quickly map between nd index and 2d index.
63 : !> \param map index mapping data
64 : !> \param dims nd sizes
65 : !> \param map1_2d which nd-indices map to first matrix index and in which order
66 : !> \param map2_2d which nd-indices map to second matrix index and in which order
67 : !> \param base base index (1 for Fortran-style, 0 for C-style, default is 1)
68 : !> \param col_major whether index should be column major order
69 : !> (.TRUE. for Fortran-style, .FALSE. for C-style, default is .TRUE.).
70 : !> \author Patrick Seewald
71 : ! **************************************************************************************************
72 13977448 : SUBROUTINE create_nd_to_2d_mapping(map, dims, map1_2d, map2_2d, base, col_major)
73 : TYPE(nd_to_2d_mapping), INTENT(OUT) :: map
74 : INTEGER, DIMENSION(:), INTENT(IN) :: dims, map1_2d, map2_2d
75 : INTEGER, INTENT(IN), OPTIONAL :: base
76 : LOGICAL, INTENT(IN), OPTIONAL :: col_major
77 :
78 : INTEGER :: i
79 :
80 3494362 : IF (PRESENT(col_major)) THEN
81 883714 : map%col_major = col_major
82 : ELSE
83 2610648 : map%col_major = .TRUE.
84 : END IF
85 :
86 3494362 : IF (PRESENT(base)) THEN
87 883714 : map%base = base
88 : ELSE
89 2610648 : map%base = 1
90 : END IF
91 :
92 3494362 : map%ndim1_2d = SIZE(map1_2d)
93 3494362 : map%ndim2_2d = SIZE(map2_2d)
94 3494362 : map%ndim_nd = SIZE(dims)
95 :
96 14960292 : ALLOCATE (map%map1_2d, source=map1_2d)
97 15027832 : ALLOCATE (map%map2_2d, source=map2_2d)
98 19505038 : ALLOCATE (map%dims_nd, source=dims)
99 15943136 : ALLOCATE (map%dims1_2d, source=dims(map1_2d))
100 16078216 : ALLOCATE (map%dims2_2d, source=dims(map2_2d))
101 :
102 10483086 : ALLOCATE (map%map_nd(map%ndim_nd))
103 15943136 : map%map_nd(map1_2d) = (/(i, i=1, SIZE(map1_2d))/)
104 16078216 : map%map_nd(map2_2d) = (/(i + SIZE(map1_2d), i=1, SIZE(map2_2d))/)
105 :
106 19505038 : map%dims_2d = [PRODUCT(INT(map%dims1_2d, KIND=int_8)), PRODUCT(INT(map%dims2_2d, KIND=int_8))]
107 :
108 3494362 : END SUBROUTINE create_nd_to_2d_mapping
109 :
110 : ! **************************************************************************************************
111 : !> \brief
112 : !> \author Patrick Seewald
113 : ! **************************************************************************************************
114 3933073 : SUBROUTINE destroy_nd_to_2d_mapping(map)
115 : TYPE(nd_to_2d_mapping), INTENT(INOUT) :: map
116 :
117 3933073 : DEALLOCATE (map%dims1_2d)
118 3933073 : DEALLOCATE (map%dims2_2d)
119 3933073 : DEALLOCATE (map%map1_2d)
120 3933073 : DEALLOCATE (map%map2_2d)
121 3933073 : DEALLOCATE (map%map_nd)
122 3933073 : DEALLOCATE (map%dims_nd)
123 3933073 : END SUBROUTINE destroy_nd_to_2d_mapping
124 :
125 : ! **************************************************************************************************
126 : !> \brief
127 : !> \author Patrick Seewald
128 : ! **************************************************************************************************
129 8258880 : PURE FUNCTION ndims_mapping(map)
130 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
131 : INTEGER :: ndims_mapping
132 :
133 8258880 : ndims_mapping = map%ndim_nd
134 8258880 : END FUNCTION
135 :
136 : ! **************************************************************************************************
137 : !> \brief how many tensor dimensions are mapped to matrix row
138 : !> \author Patrick Seewald
139 : ! **************************************************************************************************
140 15284617 : PURE FUNCTION ndims_mapping_row(map)
141 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
142 : INTEGER :: ndims_mapping_row
143 15284617 : ndims_mapping_row = map%ndim1_2d
144 15284617 : END FUNCTION
145 :
146 : ! **************************************************************************************************
147 : !> \brief how many tensor dimensions are mapped to matrix column
148 : !> \author Patrick Seewald
149 : ! **************************************************************************************************
150 15283933 : PURE FUNCTION ndims_mapping_column(map)
151 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
152 : INTEGER :: ndims_mapping_column
153 15283933 : ndims_mapping_column = map%ndim2_2d
154 15283933 : END FUNCTION
155 :
156 : ! **************************************************************************************************
157 : !> \brief get mapping info
158 : !> \param map index mapping data
159 : !> \param ndim_nd number of dimensions
160 : !> \param ndim1_2d number of dimensions that map to first 2d index
161 : !> \param ndim2_2d number of dimensions that map to first 2d index
162 : !> \param dims_2d 2d dimensions
163 : !> \param dims_nd nd dimensions
164 : !> \param dims1_2d dimensions that map to first 2d index
165 : !> \param dims2_2d dimensions that map to second 2d index
166 : !> \param map1_2d indices that map to first 2d index
167 : !> \param map2_2d indices that map to second 2d index
168 : !> \param map_nd inverse of [map1_2d, map2_2d]
169 : !> \param base base index
170 : !> \param col_major is index in column major order
171 : !> \author Patrick Seewald
172 : ! **************************************************************************************************
173 9952702 : PURE SUBROUTINE dbt_get_mapping_info(map, ndim_nd, ndim1_2d, ndim2_2d, dims_2d_i8, &
174 6146540 : dims_2d, dims_nd, dims1_2d, dims2_2d, &
175 10959774 : map1_2d, map2_2d, map_nd, base, col_major)
176 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
177 : INTEGER, INTENT(OUT), OPTIONAL :: ndim_nd, ndim1_2d, ndim2_2d
178 : INTEGER(KIND=int_8), DIMENSION(2), INTENT(OUT), OPTIONAL :: dims_2d_i8
179 : INTEGER, DIMENSION(2), INTENT(OUT), OPTIONAL :: dims_2d
180 : INTEGER, DIMENSION(ndims_mapping(map)), &
181 : INTENT(OUT), OPTIONAL :: dims_nd
182 : INTEGER, DIMENSION(ndims_mapping_row(map)), INTENT(OUT), &
183 : OPTIONAL :: dims1_2d
184 : INTEGER, DIMENSION(ndims_mapping_column(map)), INTENT(OUT), &
185 : OPTIONAL :: dims2_2d
186 : INTEGER, DIMENSION(ndims_mapping_row(map)), INTENT(OUT), &
187 : OPTIONAL :: map1_2d
188 : INTEGER, DIMENSION(ndims_mapping_column(map)), INTENT(OUT), &
189 : OPTIONAL :: map2_2d
190 : INTEGER, DIMENSION(ndims_mapping(map)), &
191 : INTENT(OUT), OPTIONAL :: map_nd
192 : INTEGER, INTENT(OUT), OPTIONAL :: base
193 : LOGICAL, INTENT(OUT), OPTIONAL :: col_major
194 :
195 9952702 : IF (PRESENT(ndim_nd)) ndim_nd = map%ndim_nd
196 9952702 : IF (PRESENT(ndim1_2d)) ndim1_2d = map%ndim1_2d
197 9952702 : IF (PRESENT(ndim2_2d)) ndim2_2d = map%ndim2_2d
198 13024832 : IF (PRESENT(dims_2d_i8)) dims_2d_i8(:) = map%dims_2d(:)
199 13476066 : IF (PRESENT(dims_2d)) dims_2d(:) = INT(map%dims_2d(:))
200 9952702 : IF (PRESENT(dims_nd)) THEN
201 14271600 : dims_nd(:) = map%dims_nd(:)
202 : END IF
203 9952702 : IF (PRESENT(dims1_2d)) THEN
204 2384569 : dims1_2d(:) = map%dims1_2d
205 : END IF
206 9952702 : IF (PRESENT(dims2_2d)) THEN
207 2583504 : dims2_2d(:) = map%dims2_2d
208 : END IF
209 9952702 : IF (PRESENT(map1_2d)) THEN
210 12758680 : map1_2d(:) = map%map1_2d
211 : END IF
212 9952702 : IF (PRESENT(map2_2d)) THEN
213 12522617 : map2_2d(:) = map%map2_2d
214 : END IF
215 9952702 : IF (PRESENT(map_nd)) THEN
216 0 : map_nd(:) = map%map_nd(:)
217 : END IF
218 9952702 : IF (PRESENT(base)) THEN
219 0 : base = map%base
220 : END IF
221 9952702 : IF (PRESENT(col_major)) THEN
222 0 : col_major = map%col_major
223 : END IF
224 :
225 27059016 : END SUBROUTINE dbt_get_mapping_info
226 :
227 : ! **************************************************************************************************
228 : !> \brief transform nd index to flat index
229 : !> \param ind_in nd index
230 : !> \param dims nd dimensions
231 : !> \param ind_out flat index
232 : !> \author Patrick Seewald
233 : ! **************************************************************************************************
234 150211962 : PURE FUNCTION combine_tensor_index(ind_in, dims) RESULT(ind_out)
235 : INTEGER, DIMENSION(:), INTENT(IN) :: ind_in, dims
236 : INTEGER(KIND=int_8) :: ind_out
237 : INTEGER :: i_dim
238 :
239 150211962 : ind_out = ind_in(SIZE(dims))
240 227635888 : DO i_dim = SIZE(dims) - 1, 1, -1
241 227635888 : ind_out = (ind_out - 1)*dims(i_dim) + ind_in(i_dim)
242 : END DO
243 :
244 150211962 : END FUNCTION
245 :
246 : ! **************************************************************************************************
247 : !> \brief transform nd index to flat index
248 : !> \param ind_in nd index
249 : !> \param dims nd dimensions
250 : !> \param ind_out flat index
251 : !> \author Patrick Seewald
252 : ! **************************************************************************************************
253 33739084 : PURE FUNCTION combine_pgrid_index(ind_in, dims) RESULT(ind_out)
254 : INTEGER, DIMENSION(:), INTENT(IN) :: ind_in, dims
255 : INTEGER :: ind_out
256 :
257 : INTEGER :: i_dim
258 :
259 33739084 : ind_out = ind_in(1)
260 48922689 : DO i_dim = 2, SIZE(dims)
261 48922689 : ind_out = ind_out*dims(i_dim) + ind_in(i_dim)
262 : END DO
263 33739084 : END FUNCTION
264 :
265 : ! **************************************************************************************************
266 : !> \brief transform flat index to nd index
267 : !> \param ind_in flat index
268 : !> \param dims nd dimensions
269 : !> \param ind_out nd index
270 : !> \author Patrick Seewald
271 : ! **************************************************************************************************
272 122802955 : PURE FUNCTION split_tensor_index(ind_in, dims) RESULT(ind_out)
273 : INTEGER(KIND=int_8), INTENT(IN) :: ind_in
274 : INTEGER, DIMENSION(:), INTENT(IN) :: dims
275 : INTEGER, DIMENSION(SIZE(dims)) :: ind_out
276 :
277 : INTEGER(KIND=int_8) :: tmp
278 : INTEGER :: i_dim
279 :
280 122802955 : tmp = ind_in
281 296014684 : DO i_dim = 1, SIZE(dims)
282 173211729 : ind_out(i_dim) = INT(MOD(tmp - 1, INT(dims(i_dim), int_8)) + 1)
283 296014684 : tmp = (tmp - 1)/dims(i_dim) + 1
284 : END DO
285 :
286 122802955 : END FUNCTION
287 :
288 : ! **************************************************************************************************
289 : !> \brief transform flat index to nd index
290 : !> \param ind_in flat index
291 : !> \param dims nd dimensions
292 : !> \param ind_out nd index
293 : !> \author Patrick Seewald
294 : ! **************************************************************************************************
295 2621772 : PURE FUNCTION split_pgrid_index(ind_in, dims) RESULT(ind_out)
296 : INTEGER, INTENT(IN) :: ind_in
297 : INTEGER, DIMENSION(:), INTENT(IN) :: dims
298 : INTEGER, DIMENSION(SIZE(dims)) :: ind_out
299 :
300 : INTEGER :: tmp
301 : INTEGER :: i_dim
302 :
303 2621772 : tmp = ind_in
304 6040064 : DO i_dim = SIZE(dims), 1, -1
305 3418292 : ind_out(i_dim) = MOD(tmp, dims(i_dim))
306 6040064 : tmp = tmp/dims(i_dim)
307 : END DO
308 2621772 : END FUNCTION
309 :
310 : ! **************************************************************************************************
311 : !> \brief transform nd index to 2d index, using info from index mapping.
312 : !> \param map index mapping
313 : !> \param ind_in nd index
314 : !> \param ind_out 2d index
315 : !> \author Patrick Seewald
316 : ! **************************************************************************************************
317 61475161 : PURE FUNCTION get_2d_indices_tensor(map, ind_in) RESULT(ind_out)
318 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
319 : INTEGER, DIMENSION(map%ndim_nd), INTENT(IN) :: ind_in
320 : INTEGER(KIND=int_8), DIMENSION(2) :: ind_out
321 : INTEGER :: i
322 : INTEGER, DIMENSION(${maxrank}$) :: ind_tmp
323 :
324 136690608 : DO i = 1, map%ndim1_2d
325 136690608 : ind_tmp(i) = ind_in(map%map1_2d(i))
326 : END DO
327 61475161 : ind_out(1) = combine_tensor_index(ind_tmp(:map%ndim1_2d), map%dims1_2d)
328 :
329 158916576 : DO i = 1, map%ndim2_2d
330 158916576 : ind_tmp(i) = ind_in(map%map2_2d(i))
331 : END DO
332 61475161 : ind_out(2) = combine_tensor_index(ind_tmp(:map%ndim2_2d), map%dims2_2d)
333 61475161 : END FUNCTION
334 :
335 : ! **************************************************************************************************
336 : !> \brief transform nd index to 2d index, using info from index mapping.
337 : !> \param map index mapping
338 : !> \param ind_in nd index
339 : !> \param ind_out 2d index
340 : !> \author Patrick Seewald
341 : ! **************************************************************************************************
342 0 : PURE FUNCTION get_2d_indices_pgrid(map, ind_in) RESULT(ind_out)
343 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
344 : INTEGER, DIMENSION(map%ndim_nd), INTENT(IN) :: ind_in
345 : INTEGER, DIMENSION(2) :: ind_out
346 : INTEGER :: i
347 : INTEGER, DIMENSION(${maxrank}$) :: ind_tmp
348 :
349 0 : DO i = 1, map%ndim1_2d
350 0 : ind_tmp(i) = ind_in(map%map1_2d(i))
351 : END DO
352 0 : ind_out(1) = combine_pgrid_index(ind_tmp(:map%ndim1_2d), map%dims1_2d)
353 :
354 0 : DO i = 1, map%ndim2_2d
355 0 : ind_tmp(i) = ind_in(map%map2_2d(i))
356 : END DO
357 0 : ind_out(2) = combine_pgrid_index(ind_tmp(:map%ndim2_2d), map%dims2_2d)
358 0 : END FUNCTION
359 :
360 : ! **************************************************************************************************
361 : !> \brief transform 2d index to nd index, using info from index mapping.
362 : !> \param map index mapping
363 : !> \param ind_in 2d index
364 : !> \param ind_out nd index
365 : !> \author Patrick Seewald
366 : ! **************************************************************************************************
367 36684840 : PURE FUNCTION get_nd_indices_tensor(map, ind_in) RESULT(ind_out)
368 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
369 : INTEGER(KIND=int_8), DIMENSION(2), INTENT(IN) :: ind_in
370 : INTEGER, DIMENSION(map%ndim_nd) :: ind_out
371 : INTEGER, DIMENSION(${maxrank}$) :: ind_tmp
372 : INTEGER :: i
373 :
374 36684840 : ind_tmp(:map%ndim1_2d) = split_tensor_index(ind_in(1), map%dims1_2d)
375 :
376 86636720 : DO i = 1, map%ndim1_2d
377 86636720 : ind_out(map%map1_2d(i)) = ind_tmp(i)
378 : END DO
379 :
380 36684840 : ind_tmp(:map%ndim2_2d) = split_tensor_index(ind_in(2), map%dims2_2d)
381 :
382 90459547 : DO i = 1, map%ndim2_2d
383 90459547 : ind_out(map%map2_2d(i)) = ind_tmp(i)
384 : END DO
385 :
386 36684840 : END FUNCTION
387 :
388 : ! **************************************************************************************************
389 : !> \brief transform 2d index to nd index, using info from index mapping.
390 : !> \param map index mapping
391 : !> \param ind_in 2d index
392 : !> \param ind_out nd index
393 : !> \author Patrick Seewald
394 : ! **************************************************************************************************
395 851752 : PURE FUNCTION get_nd_indices_pgrid(map, ind_in) RESULT(ind_out)
396 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map
397 : INTEGER, DIMENSION(2), INTENT(IN) :: ind_in
398 : INTEGER, DIMENSION(map%ndim_nd) :: ind_out
399 :
400 3817656 : ind_out(map%map1_2d) = split_pgrid_index(ind_in(1), map%dims1_2d)
401 3927168 : ind_out(map%map2_2d) = split_pgrid_index(ind_in(2), map%dims2_2d)
402 :
403 851752 : END FUNCTION
404 :
405 : ! **************************************************************************************************
406 : !> \brief Invert order
407 : !> \author Patrick Seewald
408 : ! **************************************************************************************************
409 3240408 : PURE FUNCTION dbt_inverse_order(order)
410 : INTEGER, DIMENSION(:), INTENT(IN) :: order
411 : INTEGER, DIMENSION(SIZE(order)) :: dbt_inverse_order
412 :
413 : INTEGER :: i
414 :
415 23795356 : dbt_inverse_order(order) = (/(i, i=1, SIZE(order))/)
416 : END FUNCTION
417 :
418 : ! **************************************************************************************************
419 : !> \brief reorder tensor index (no data)
420 : !> \author Patrick Seewald
421 : ! **************************************************************************************************
422 6931296 : SUBROUTINE permute_index(map_in, map_out, order)
423 : TYPE(nd_to_2d_mapping), INTENT(IN) :: map_in
424 : TYPE(nd_to_2d_mapping), INTENT(OUT) :: map_out
425 : INTEGER, DIMENSION(ndims_mapping(map_in)), &
426 : INTENT(IN) :: order
427 :
428 : INTEGER :: ndim_nd
429 3465648 : INTEGER, DIMENSION(ndims_mapping_row(map_in)) :: map1_2d, map1_2d_reorder
430 3465648 : INTEGER, DIMENSION(ndims_mapping_column(map_in)) :: map2_2d, map2_2d_reorder
431 1732824 : INTEGER, DIMENSION(ndims_mapping(map_in)) :: dims_nd, dims_reorder
432 :
433 1732824 : CALL dbt_get_mapping_info(map_in, ndim_nd, dims_nd=dims_nd, map1_2d=map1_2d, map2_2d=map2_2d)
434 :
435 6353028 : dims_reorder(order) = dims_nd
436 :
437 4170354 : map1_2d_reorder(:) = order(map1_2d)
438 3915498 : map2_2d_reorder(:) = order(map2_2d)
439 :
440 1732824 : CALL create_nd_to_2d_mapping(map_out, dims_reorder, map1_2d_reorder, map2_2d_reorder)
441 1732824 : END SUBROUTINE
442 :
443 0 : END MODULE dbt_index
|
