group Subroutine

  subroutine group(array, array_groups, tran_group_threshold, error, comm)
    !================================================!

    use w90_constants, only: dp
    use w90_error, only: w90_error_type, set_error_alloc, set_error_dealloc

    implicit none

    ! arguments
    type(w90_error_type), allocatable, intent(out) :: error
    type(w90_comm_type), intent(in) :: comm
    real(kind=dp), intent(in) :: tran_group_threshold
    real(kind=dp), intent(in) :: array(:, :)
    integer, intent(out), allocatable :: array_groups(:)

    ! local variables
    integer, allocatable :: dummy_array(:)
    integer :: array_idx, i, j, group_number, array_size, ierr
    logical, allocatable :: logic(:)

    array_size = size(array, 2)

    allocate (dummy_array(array_size), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating dummy_array in group', comm)
      return
    end if
    allocate (logic(array_size), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating logic in group', comm)
      return
    end if

    !Initialise dummy array
    dummy_array = 0

    !Initialise logic to false
    logic = .false.

    !Define counter of number of groups
    array_idx = 1

    !Loop over columns of array (ie array_size)
    do i = 1, array_size

      !If an element of logic is true then it means the wannier function has already been grouped
      if (logic(i) .eqv. .false.) then

        !Create a group for the wannier function
        logic(i) = .true.

        !Initialise the number of wannier functions in this group to be 1
        group_number = 1

        !Loop over the rest of wannier functions in array
        do j = min(i + 1, array_size), array_size

          !Special termination cases
          if ((j .eq. 1) .or. (i .eq. array_size)) then
            dummy_array(array_idx) = group_number
            exit
          end if
          if (j .eq. array_size .and. (abs(array(2, j) - array(2, i)) .le. tran_group_threshold)) then
            group_number = group_number + 1
            dummy_array(array_idx) = group_number
            logic(j) = .true.
            exit
          end if

          !Check distance between wannier function_i and wannier function_j
          if (abs(array(2, j) - array(2, i)) .le. tran_group_threshold) then

            !Increment number of wannier functions in group

            group_number = group_number + 1

            !Assigns wannier function to the group

            logic(j) = .true.
          else

            !Group is finished and store number of wanniers in the group to dummy_array
            dummy_array(array_idx) = group_number

            !Increment number of groups
            array_idx = array_idx + 1
            exit
          end if
        end do
      end if
    end do

    !Copy elements of dummy_array to array_groups

    allocate (array_groups(array_idx), stat=ierr)
    if (ierr /= 0) then
      call set_error_alloc(error, 'Error in allocating array_groups in group', comm)
      return
    end if
    array_groups = dummy_array(:array_idx)

    deallocate (dummy_array, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating dummy_array in group', comm)
      return
    end if
    deallocate (logic, stat=ierr)
    if (ierr /= 0) then
      call set_error_dealloc(error, 'Error deallocating logic in group', comm)
      return
    end if

    return

  end subroutine group