split_5Drdata Subroutine

  subroutine split_5Drdata( &
       data, left, right, dim, &
       left_size, right_size, &
       shuffle, seed &
  )
    !! Split a 5D array along one dimension
    implicit none

    ! Arguments
    real(real32), dimension(:,:,:,:,:), intent(in) :: data
    !! 5D array to be split
    real(real32), allocatable, dimension(:,:,:,:,:), intent(out) :: left, right
    !! 5D arrays to store the left and right splits
    integer, intent(in) :: dim
    !! Dimension along which to split
    real(real32), optional, intent(in) :: left_size, right_size
    !! Size of the left and right splits
    logical, optional, intent(in) :: shuffle
    !! Shuffle the data before splitting
    integer, optional, intent(in) :: seed
    !! Random seed

    ! Local variables
    integer :: seed_, left_num_, right_num_
    !! Random seed, number of elements in left and right splits
    logical :: shuffle_
    !! Shuffle flag
    integer :: i, j
    !! Loop indices
    integer :: num_redos
    !! Number of redos
    real(real32) :: rtmp1
    !! Temporary real
    integer, allocatable, dimension(:) :: indices_l, indices_r
    !! Index arrays
    real(real32), allocatable, dimension(:) :: tlist
    !! Temporary list
    real(real32), allocatable, dimension(:,:,:,:,:) :: data_copy
    !! Copy of the input data

    type :: idx_type
       !! Type for index array
       integer, allocatable, dimension(:) :: loc
       !! Index array
    end type idx_type
    type(idx_type), dimension(5) :: idx
    !! Index array


    ! Determine number of elements for left and right split
    !---------------------------------------------------------------------------
    if(.not.present(left_size).and..not.present(right_size))then
       call stop_program("neither left_size nor right_size provided to split. &
            &Expected at least one." &
       )
       return
    elseif(present(left_size).and..not.present(right_size))then
       left_num_  = nint(left_size*size(data,dim))
       right_num_ = size(data,dim) - left_num_
    elseif(.not.present(left_size).and.present(right_size))then
       right_num_ = nint(right_size*size(data,dim))
       left_num_  = size(data,dim) - right_num_
    else
       left_num_  = nint(left_size*size(data,dim))
       right_num_ = nint(right_size*size(data,dim))
       if(left_num_ + right_num_ .ne. size(data,dim)) &
            right_num_ = size(data,dim) - left_num_
    end if


    ! Initialies optional arguments
    !---------------------------------------------------------------------------
    if(present(shuffle))then
       shuffle_ = shuffle
    else
       shuffle_ = .false.
    end if

    if(present(seed))then
       seed_ = seed
    else
       call system_clock(count=seed_)
    end if


    ! Copy input data
    !---------------------------------------------------------------------------
    data_copy = data
    if(shuffle_) call shuffle_5Drdata(data_copy,dim,seed_)


    ! Get list of indices for right split
    !---------------------------------------------------------------------------
    num_redos = 0
    allocate(tlist(right_num_))
    call random_number(tlist)
    indices_r = floor(tlist*size(data,dim)) + 1
    i = 1
    indices_r_loop: do
       if(i.ge.right_num_) exit indices_r_loop
       i = i + 1
       if(any(indices_r(:i-1).eq.indices_r(i)))then
          indices_r(i:right_num_-num_redos-1) = &
               indices_r(i+1:right_num_-num_redos)
          call random_number(rtmp1)
          indices_r(right_num_) = floor(rtmp1*size(data,dim)) + 1
          i = i - 1
       end if
    end do indices_r_loop


    ! Generate right split
    !---------------------------------------------------------------------------
    do i=1,5
       if(i.eq.dim)then
          idx(i)%loc = indices_r
       else
          idx(i)%loc = (/ ( j, j=1,size(data,i) ) /)
       end if
    end do
    right = data_copy(idx(1)%loc,idx(2)%loc,idx(3)%loc,idx(4)%loc,idx(5)%loc)


    ! Get list of indices for left split
    !---------------------------------------------------------------------------
    indices_l_loop: do i=1,size(data,dim)
       if(any(indices_r.eq.i)) cycle indices_l_loop
       if(allocated(indices_l))then
          indices_l = [indices_l(:), i]
       else
          indices_l = [i]
       end if
    end do indices_l_loop


    ! Generate left split
    !---------------------------------------------------------------------------
    idx(dim)%loc = indices_l
    left = data_copy(idx(1)%loc,idx(2)%loc,idx(3)%loc,idx(4)%loc,idx(5)%loc)

  end subroutine split_5Drdata