Line data    Source code

       1             : !--------------------------------------------------------------------------------------------------!
       2             : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3             : !   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
       4             : !                                                                                                  !
       5             : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6             : !--------------------------------------------------------------------------------------------------!
       7             : 
       8             : ! **************************************************************************************************
       9             : !> \brief Representation of arbitrary number of 1d integer arrays with arbitrary sizes.
      10             : !>        This is needed for generic handling of dimension-specific tensor quantities
      11             : !>        (such as block index).
      12             : !> \author Patrick Seewald
      13             : ! **************************************************************************************************
      14             : MODULE dbt_array_list_methods
      15             : 
      16             :    #:include "dbt_macros.fypp"
      17             :    #:set maxdim = maxrank
      18             :    #:set ndims = range(2,maxdim+1)
      19             : 
      20             :    USE dbt_index, ONLY: dbt_inverse_order
      21             :    USE dbt_allocate_wrap, ONLY: allocate_any
      22             : 
      23             : #include "../base/base_uses.f90"
      24             :    IMPLICIT NONE
      25             :    PRIVATE
      26             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'dbt_array_list_methods'
      27             : 
      28             :    PUBLIC  :: &
      29             :       array_eq_i, &
      30             :       array_list, &
      31             :       array_offsets, &
      32             :       array_sublist, &
      33             :       create_array_list, &
      34             :       destroy_array_list, &
      35             :       get_array_elements, &
      36             :       get_arrays, &
      37             :       get_ith_array, &
      38             :       number_of_arrays, &
      39             :       reorder_arrays, &
      40             :       sizes_of_arrays, &
      41             :       sum_of_arrays, &
      42             :       check_equal
      43             : 
      44             :    TYPE array_list
      45             :       INTEGER, DIMENSION(:), ALLOCATABLE :: col_data
      46             :       INTEGER, DIMENSION(:), ALLOCATABLE :: ptr
      47             :    END TYPE
      48             : 
      49             :    INTERFACE get_ith_array
      50             :       MODULE PROCEDURE allocate_and_get_ith_array
      51             :       MODULE PROCEDURE get_ith_array
      52             :    END INTERFACE
      53             : 
      54             : CONTAINS
      55             : 
      56             : ! **************************************************************************************************
      57             : !> \brief number of arrays stored in list
      58             : !> \author Patrick Seewald
      59             : ! **************************************************************************************************
      60   143589516 :    PURE FUNCTION number_of_arrays(list)
      61             :       TYPE(array_list), INTENT(IN) :: list
      62             :       INTEGER                      :: number_of_arrays
      63             : 
      64   143589516 :       number_of_arrays = SIZE(list%ptr) - 1
      65             : 
      66   143589516 :    END FUNCTION number_of_arrays
      67             : 
      68             : ! **************************************************************************************************
      69             : !> \brief Get an element for each array.
      70             : !> \param indices element index for each array
      71             : !> \author Patrick Seewald
      72             : ! **************************************************************************************************
      73   123034462 :    PURE FUNCTION get_array_elements(list, indices)
      74             :       TYPE(array_list), INTENT(IN)                           :: list
      75             :       INTEGER, DIMENSION(number_of_arrays(list)), INTENT(IN) :: indices
      76             :       INTEGER, DIMENSION(number_of_arrays(list))             :: get_array_elements
      77             : 
      78             :       INTEGER                                                :: i, ind
      79             : 
      80   386551059 :       DO i = 1, SIZE(indices)
      81   263516597 :          ind = indices(i) + list%ptr(i) - 1
      82   386551059 :          get_array_elements(i) = list%col_data(ind)
      83             :       END DO
      84             : 
      85             :    END FUNCTION get_array_elements
      86             : 
      87             : ! **************************************************************************************************
      88             : !> \brief collects any number of arrays of different sizes into a single array (list%col_data),
      89             : !>        storing the indices that start a new array (list%ptr).
      90             : !> \param list list of arrays
      91             : !> \param ndata number of arrays
      92             : !> \param data arrays 1 and 2
      93             : !> \author Patrick Seewald
      94             : ! **************************************************************************************************
      95     5281202 :    SUBROUTINE create_array_list(list, ndata, ${varlist("data")}$)
      96             :       TYPE(array_list), INTENT(OUT)               :: list
      97             :       INTEGER, INTENT(IN)                         :: ndata
      98             :       INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: ${varlist("data")}$
      99             :       INTEGER                                     :: ptr, size_all
     100             : 
     101     5281202 :       size_all = 0
     102             : 
     103             :       #:for dim in range(1, maxdim+1)
     104    16485898 :          IF (ndata .GE. ${dim}$) THEN
     105    11206552 :             CPASSERT(PRESENT(data_${dim}$))
     106    11206552 :             size_all = size_all + SIZE(data_${dim}$)
     107             :          END IF
     108             :       #:endfor
     109             : 
     110    15843606 :       ALLOCATE (list%ptr(ndata + 1))
     111    15733395 :       ALLOCATE (list%col_data(size_all))
     112             : 
     113     5281202 :       ptr = 1
     114     5281202 :       list%ptr(1) = ptr
     115             : 
     116             :       #:for dim in range(1, maxdim+1)
     117    16485898 :          IF (ndata .GE. ${dim}$) THEN
     118   125567011 :             list%col_data(ptr:ptr + SIZE(data_${dim}$) - 1) = data_${dim}$ (:)
     119    11206552 :             ptr = ptr + SIZE(data_${dim}$)
     120    11206552 :             list%ptr(${dim+1}$) = ptr
     121             :          END IF
     122             :       #:endfor
     123             : 
     124     5281202 :    END SUBROUTINE
     125             : 
     126             : ! **************************************************************************************************
     127             : !> \brief extract a subset of arrays
     128             : !> \param list list of arrays
     129             : !> \param i_selected array numbers to retrieve
     130             : !> \author Patrick Seewald
     131             : ! **************************************************************************************************
     132     2227972 :    FUNCTION array_sublist(list, i_selected)
     133             :       TYPE(array_list), INTENT(IN)                           :: list
     134             :       INTEGER, DIMENSION(:), INTENT(IN)                      :: i_selected
     135             :       TYPE(array_list)                                       :: array_sublist
     136             :       INTEGER :: ndata
     137     1113986 :       INTEGER, ALLOCATABLE, DIMENSION(:) :: ${varlist("data")}$
     138             : 
     139     1113986 :       ndata = SIZE(i_selected)
     140             : 
     141             :       #:for dim in range(1, maxdim+1)
     142     2227972 :          IF (ndata == ${dim}$) THEN
     143     1113986 :             CALL get_arrays(list, ${varlist("data", nmax=dim)}$, i_selected=i_selected)
     144     1113986 :             CALL create_array_list(array_sublist, ndata, ${varlist("data", nmax=dim)}$)
     145             :          END IF
     146             :       #:endfor
     147     2227972 :    END FUNCTION
     148             : 
     149             : ! **************************************************************************************************
     150             : !> \brief destroy array list.
     151             : !> \author Patrick Seewald
     152             : ! **************************************************************************************************
     153     4105464 :    SUBROUTINE destroy_array_list(list)
     154             :       TYPE(array_list), INTENT(INOUT) :: list
     155             : 
     156     4105464 :       DEALLOCATE (list%ptr, list%col_data)
     157     4105464 :    END SUBROUTINE
     158             : 
     159             : ! **************************************************************************************************
     160             : !> \brief Get all arrays contained in list
     161             : !> \param data arrays 1 and 2
     162             : !> \param i_selected array numbers to retrieve (if not present, all arrays are returned)
     163             : !> \author Patrick Seewald
     164             : ! **************************************************************************************************
     165     4234393 :    SUBROUTINE get_arrays(list, ${varlist("data")}$, i_selected)
     166             :       !! Get all arrays contained in list
     167             :       TYPE(array_list), INTENT(IN)                       :: list
     168             :       INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT), &
     169             :          OPTIONAL                                        :: ${varlist("data")}$
     170             :       INTEGER, DIMENSION(:), INTENT(IN), &
     171             :          OPTIONAL                                        :: i_selected
     172             :       INTEGER                                            :: i, ndata
     173     4234393 :       INTEGER, DIMENSION(number_of_arrays(list))         :: o
     174             : 
     175    14179159 :       o(:) = 0
     176     4234393 :       IF (PRESENT(i_selected)) THEN
     177     3094815 :          ndata = SIZE(i_selected)
     178     9752053 :          o(1:ndata) = i_selected(:)
     179             :       ELSE
     180     1139578 :          ndata = number_of_arrays(list)
     181     4972514 :          o(1:ndata) = (/(i, i=1, ndata)/)
     182             :       END IF
     183             : 
     184             :       ASSOCIATE (ptr => list%ptr, col_data => list%col_data)
     185             :          #:for dim in range(1, maxdim+1)
     186    12807439 :             IF (ndata > ${dim-1}$) THEN
     187   102301709 :                ALLOCATE (data_${dim}$, source=col_data(ptr(o(${dim}$)):ptr(o(${dim}$) + 1) - 1))
     188             :             END IF
     189             :          #:endfor
     190             :       END ASSOCIATE
     191             : 
     192     4234393 :    END SUBROUTINE get_arrays
     193             : 
     194             : ! **************************************************************************************************
     195             : !> \brief get ith array
     196             : !> \author Patrick Seewald
     197             : ! **************************************************************************************************
     198     1135442 :    SUBROUTINE get_ith_array(list, i, array_size, array)
     199             :       TYPE(array_list), INTENT(IN)                    :: list
     200             :       INTEGER, INTENT(IN)                             :: i
     201             :       INTEGER, INTENT(IN)                             :: array_size
     202             :       INTEGER, DIMENSION(array_size), INTENT(OUT)     :: array
     203             : 
     204             :       ASSOCIATE (ptr => list%ptr, col_data => list%col_data)
     205           0 :          CPASSERT(i <= number_of_arrays(list))
     206             : 
     207     5926729 :          array(:) = col_data(ptr(i):ptr(i + 1) - 1)
     208             : 
     209             :       END ASSOCIATE
     210             : 
     211     1135442 :    END SUBROUTINE
     212             : 
     213             : ! **************************************************************************************************
     214             : !> \brief get ith array
     215             : !> \author Patrick Seewald
     216             : ! **************************************************************************************************
     217     1496818 :    SUBROUTINE allocate_and_get_ith_array(list, i, array)
     218             :       TYPE(array_list), INTENT(IN)                    :: list
     219             :       INTEGER, INTENT(IN)                             :: i
     220             :       INTEGER, DIMENSION(:), ALLOCATABLE, INTENT(OUT) :: array
     221             : 
     222             :       ASSOCIATE (ptr => list%ptr, col_data => list%col_data)
     223           0 :          CPASSERT(i <= number_of_arrays(list))
     224             : 
     225    15222138 :          ALLOCATE (array, source=col_data(ptr(i):ptr(i + 1) - 1))
     226             :       END ASSOCIATE
     227     1496818 :    END SUBROUTINE
     228             : 
     229             : ! **************************************************************************************************
     230             : !> \brief sizes of arrays stored in list
     231             : !> \author Patrick Seewald
     232             : ! **************************************************************************************************
     233     3803308 :    FUNCTION sizes_of_arrays(list)
     234             :       TYPE(array_list), INTENT(IN)       :: list
     235             :       INTEGER, ALLOCATABLE, DIMENSION(:) :: sizes_of_arrays
     236             : 
     237             :       INTEGER                            :: i_data, num_data
     238             : 
     239     1901654 :       num_data = number_of_arrays(list)
     240     5704962 :       ALLOCATE (sizes_of_arrays(num_data))
     241     5762230 :       DO i_data = 1, num_data
     242     5762230 :          sizes_of_arrays(i_data) = list%ptr(i_data + 1) - list%ptr(i_data)
     243             :       END DO
     244             :    END FUNCTION sizes_of_arrays
     245             : 
     246             : ! **************************************************************************************************
     247             : !> \brief sum of all elements for each array stored in list
     248             : !> \author Patrick Seewald
     249             : ! **************************************************************************************************
     250     1114536 :    FUNCTION sum_of_arrays(list)
     251             :       TYPE(array_list), INTENT(IN)       :: list
     252             :       INTEGER, ALLOCATABLE, DIMENSION(:) :: sum_of_arrays
     253             : 
     254             :       INTEGER                            :: i_data, num_data
     255             : 
     256      557268 :       num_data = number_of_arrays(list)
     257     1671804 :       ALLOCATE (sum_of_arrays(num_data))
     258     1471936 :       DO i_data = 1, num_data
     259     9642960 :          sum_of_arrays(i_data) = SUM(list%col_data(list%ptr(i_data):list%ptr(i_data + 1) - 1))
     260             :       END DO
     261             : 
     262             :    END FUNCTION sum_of_arrays
     263             : 
     264             : ! **************************************************************************************************
     265             : !> \brief partial sums of array elements.
     266             : !> \author Patrick Seewald
     267             : ! **************************************************************************************************
     268      185756 :    SUBROUTINE array_offsets(list_in, list_out)
     269             :       TYPE(array_list), INTENT(IN)  :: list_in
     270             :       TYPE(array_list), INTENT(OUT) :: list_out
     271             : 
     272             :       INTEGER                       :: i_data, i_ptr, num_data, partial_sum
     273             : 
     274      185756 :       num_data = number_of_arrays(list_in)
     275     1200358 :       ALLOCATE (list_out%ptr, source=list_in%ptr)
     276      557268 :       ALLOCATE (list_out%col_data(SIZE(list_in%col_data)))
     277      643090 :       DO i_data = 1, num_data
     278      457334 :          partial_sum = 1
     279     4728602 :          DO i_ptr = list_out%ptr(i_data), list_out%ptr(i_data + 1) - 1
     280     4085512 :             list_out%col_data(i_ptr) = partial_sum
     281     4542846 :             partial_sum = partial_sum + list_in%col_data(i_ptr)
     282             :          END DO
     283             :       END DO
     284      185756 :    END SUBROUTINE
     285             : 
     286             : ! **************************************************************************************************
     287             : !> \brief reorder array list.
     288             : !> \author Patrick Seewald
     289             : ! **************************************************************************************************
     290     1980829 :    SUBROUTINE reorder_arrays(list_in, list_out, order)
     291             :       TYPE(array_list), INTENT(IN)                     :: list_in
     292             :       TYPE(array_list), INTENT(OUT)                    :: list_out
     293     1980829 :       INTEGER, ALLOCATABLE, DIMENSION(:)               :: ${varlist("data")}$
     294             :       INTEGER, DIMENSION(number_of_arrays(list_in)), &
     295             :          INTENT(IN)                                    :: order
     296             : 
     297             :       #:for ndim in ndims
     298     3961610 :          IF (number_of_arrays(list_in) == ${ndim}$) THEN
     299     1980829 :             CALL get_arrays(list_in, ${varlist("data", nmax=ndim)}$, i_selected=dbt_inverse_order(order))
     300             :             CALL create_array_list(list_out, number_of_arrays(list_in), &
     301     1980829 :                                    ${varlist("data", nmax=ndim)}$)
     302             :          END IF
     303             :       #:endfor
     304             : 
     305     1980829 :    END SUBROUTINE
     306             : 
     307             : ! **************************************************************************************************
     308             : !> \brief check whether two array lists are equal
     309             : !> \author Patrick Seewald
     310             : ! **************************************************************************************************
     311      621478 :    FUNCTION check_equal(list1, list2)
     312             :       TYPE(array_list), INTENT(IN)  :: list1, list2
     313             :       LOGICAL :: check_equal
     314             : 
     315      621478 :       check_equal = array_eq_i(list1%col_data, list2%col_data) .AND. array_eq_i(list1%ptr, list2%ptr)
     316      621478 :    END FUNCTION
     317             : 
     318             : ! **************************************************************************************************
     319             : !> \brief check whether two arrays are equal
     320             : !> \author Patrick Seewald
     321             : ! **************************************************************************************************
     322     2930055 :    PURE FUNCTION array_eq_i(arr1, arr2)
     323             :       INTEGER, INTENT(IN), DIMENSION(:) :: arr1
     324             :       INTEGER, INTENT(IN), DIMENSION(:) :: arr2
     325             :       LOGICAL                           :: array_eq_i
     326             : 
     327     2930055 :       array_eq_i = .FALSE.
     328    23745561 :       IF (SIZE(arr1) .EQ. SIZE(arr2)) array_eq_i = ALL(arr1 == arr2)
     329             : 
     330     2930055 :    END FUNCTION
     331             : 
     332           0 : END MODULE dbt_array_list_methods