get_sample_array Module Function

  module function get_sample_array( &
       input, start_index, end_index, batch_size, as_graph &
  ) result(sample)
    !! Get samples of batch size from a derived type array
    implicit none

    ! Arguments
    integer, intent(in) :: start_index, end_index
    !! Start and end indices
    integer, intent(in) :: batch_size
    !! Batch size
    class(array_type), dimension(:,:), intent(in) :: input
    !! Input array
    logical, intent(in) :: as_graph
    !! Boolean whether to treat the input as a graph

    type(array_type), dimension(:,:), allocatable :: sample
    !! Sample array

    ! Local variables
    integer :: i, j
    !! Loop index

    if(as_graph)then
       allocate(sample(size(input,1), batch_size))
       do i = 1, size(input,1)
          do j = start_index, end_index, 1
             sample(i, j - start_index + 1)%val = input(i, j)%val
          end do
       end do
    else
       allocate(sample(size(input,1), size(input,2)))
       do i = 1, size(input,1)
          do j = 1, size(input,2)
             call sample(i,j)%allocate(array_shape=[input(i,j)%shape, &
                  end_index - start_index + 1])
             sample(i,j)%val = input(i,j)%val(:,start_index:end_index)
          end do
       end do
    end if

  end function get_sample_array