module kzgrid

  ! Handles the conversion of indices returned by the indexer into 
  ! ak (=wavenumber) and z (=Gabor positons).

  implicit none
  integer, parameter, private :: r8 = selected_real_kind(12,100)
  real(r8), parameter, private :: twopi = 6.2831853071795865_r8
  character(*), parameter, private :: version = &
       & '$Id: xkzgrid.f90,v 1.5 2003/10/01 19:13:38 pletzer Exp $'

  type kzgrid_obj
     integer ndims
     integer, pointer :: L(:), J(:)
     real(r8), pointer :: xmin(:), xmax(:)
     real(r8), pointer :: delta_k(:), delta_z(:)
     real(r8), pointer :: sampling(:)
     integer, pointer :: map_kz_indices(:)
     
  end type kzgrid_obj

  contains 

    subroutine kzgrid_init(this, xmin, xmax, L, J, sampling, map_kz_indices, ier)
      
      ! Constructor

      implicit none
      type(kzgrid_obj) :: this
      ! min/max points in the domain
      real(r8), intent(in) :: xmin(:), xmax(:)      
      integer, intent(in) :: L(:) ! L+1 envelopes
      integer, intent(in) :: J(:) ! 2*J+1 Fourier modes
      real(r8), intent(in) :: sampling(:) ! sampling rate should be > 1

      ! map_kz_indices maps the set of indices returned by the indexer into
      ! the indices of k and z (in that order). This provides the flexibility
      ! the reorder the memory layout for best performance. 
      !
      ! Example: in 3D, a possible layout could be map_kz_indices = (/2,3,4, 5,6,7/),
      ! ie the 2nd, 3rd, and 4th indices give the k-vector while indices 5, 6, and 7 
      ! yield the Gabor positions. In this case index 1 would be associated
      ! with field components, for instance. If one rather prefers to iterate first over 
      ! the k-vector and then over field components, set map_kz_indices = (/1,2,3, 5,6,7/). 
      integer, intent(in) :: map_kz_indices(:)
      integer, intent(out) :: ier

      ier = 0
      this % ndims = size(xmin)
      allocate(this % xmin(this % ndims))
      allocate(this % xmax(this % ndims))
      allocate(this % delta_k(this % ndims))
      allocate(this % delta_z(this % ndims))
      allocate(this % L(this % ndims))
      allocate(this % J(this % ndims))
      allocate(this % sampling(this % ndims))
      allocate(this % map_kz_indices(2*this % ndims))

      this % xmin = xmin
      this % xmax = xmax
      this % sampling = sampling
      this % L = L
      this % J = J
      this % map_kz_indices = map_kz_indices
      
      this % delta_z = (xmax - xmin)/L
      this % delta_k = twopi/(this%delta_z * sampling)

    end subroutine kzgrid_init

    subroutine kzgrid_free(this, ier)

      ! Destructor

      implicit none
      type(kzgrid_obj) :: this
      integer, intent(out) :: ier

      ier = 0
      deallocate(this % xmin)
      deallocate(this % xmax)
      deallocate(this % delta_k)
      deallocate(this % delta_z)
      deallocate(this % L)
      deallocate(this % J)
      deallocate(this % sampling)
      deallocate(this % map_kz_indices)
      
    end subroutine kzgrid_free

    subroutine kzgrid_get(this, indices, ak, z, ier)

      ! Given a set of indices, return the components of 
      ! x, k, and z.

      implicit none
      type(kzgrid_obj) :: this
      integer, intent(in) :: indices(:) ! from indexer
      real(r8), intent(out) :: ak(:), z(:) ! Wavenumber and Gabor position
      integer, intent(out) :: ier

      integer j

      ier = 0
      do j = 1, this % ndims
         ak(j) = indices(this%map_kz_indices(j          )) * this%delta_k(j)
         z(j)  = this%xmin(j) +  &
              &  indices(this%map_kz_indices(j+this%ndims)) * this%delta_z(j)
      enddo
      
    end subroutine kzgrid_get

    subroutine kzgrid_test(this, ier)

      ! Tester

      implicit none
      type(kzgrid_obj) :: this
      integer, intent(out) :: ier

      integer, parameter :: ndims = 2 ! 2-d space
      integer, parameter, dimension(ndims) :: L = (/5, 8/) ! no of Gabor intervals
      integer, parameter, dimension(ndims) :: J = (/2, 1/) ! no of positive *and* negative Fourier modes
      ! computational domain
      real(r8), parameter, dimension(ndims) :: xmin = (/0._r8, -1._r8/)
      real(r8), parameter, dimension(ndims) :: xmax = (/1._r8, +1._r8/)
      real(r8), parameter, dimension(ndims) :: sampling = (/2.0_r8, 4.0_r8/) ! sampling rate > 1
      integer :: map_kz_indices(2*ndims), indices(2*ndims+1)
      real(r8) :: ak(ndims), z(ndims)

      ier = 0

      print *,'testing kzgrid ', trim(version)
      ! kx, ky correspond to indices 2 and 3
      ! zx, zy correspond to indices 4 and 5
      map_kz_indices = (/2,3, 4,5/)
      print *,'ndims=', ndims
      print *,'L=', L, ' J=',J
      print '(2(a,2f10.3))', 'xmin=', xmin, ' xmax=', xmax
      print '(a,2f10.3)','sampling rate=', sampling
      print *,'index mapping=', map_kz_indices 

      call kzgrid_init(this, xmin, xmax, L, J, sampling, map_kz_indices, ier)

      indices = (/1, -2, -1, 0, 8/)
      call kzgrid_get(this, indices, ak, z, ier)
      print *,'indices=', indices
      print '(2(a,2f10.3))',' ak=', ak, ' z=',z

      indices = (/2, +2, +1, 0, 8/)
      call kzgrid_get(this, indices, ak, z, ier)
      print *,'indices=', indices
      print '(2(a,2f10.3))',' ak=', ak, ' z=',z
      
      indices = (/2, 0, 0, 0, 8/)
      call kzgrid_get(this, indices, ak, z, ier)
      print *,'indices=', indices
      print '(2(a,2f10.3))',' ak=', ak, ' z=',z

      indices = (/2, 1, 0, 0, 8/)
      call kzgrid_get(this, indices, ak, z, ier)
      print *,'indices=', indices
      print '(2(a,2f10.3))',' ak=', ak, ' z=',z

      indices = (/2, 1, 0, 5, 0/)
      call kzgrid_get(this, indices, ak, z, ier)
      print *,'indices=', indices
      print '(2(a,2f10.3))',' ak=', ak, ' z=',z

      call kzgrid_free(this, ier)

    end subroutine kzgrid_test

end module kzgrid

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

module xgrid

  ! Handles the conversion of indices returned by the indexer into 
  ! x (= collocation) positions.

  implicit none
  integer, parameter, private :: r8 = selected_real_kind(12,100)
  real(r8), parameter, private :: twopi = 6.2831853071795865_r8
  character(*), parameter, private :: version = &
       & '$Id: xkzgrid.f90,v 1.5 2003/10/01 19:13:38 pletzer Exp $'

  type xgrid_obj
     integer ndims
     integer, pointer :: M(:) ! no of intervals
     real(r8), pointer :: xmin(:), xmax(:)
     real(r8), pointer :: delta_x(:)
     integer, pointer :: map_x_indices(:)     
  end type xgrid_obj

  contains 

    subroutine xgrid_init(this, xmin, xmax, M, map_x_indices, ier)
      
      ! Constructor

      implicit none
      type(xgrid_obj) :: this
      ! min/max points in the domain
      real(r8), intent(in) :: xmin(:), xmax(:)      
      integer, intent(in) :: M(:) ! M+1 collocation points

      ! map_x_indices maps the set of indices returned by the indexer into
      ! the indices of collocation points. This provides the flexibility
      ! the reorder the memory layout for best performance. 
      !
      ! Example: in 3D, a possible layout could be map_x_indices = (/2,3,4/),
      ! ie the 2nd, 3rd and 4th indices give the x-point. In this case index 1  
      ! would be associated with field components, for instance. If one rather prefers 
      ! to iterate first over x-components and then over field components, set
      ! map_x_indices = (/1,2,3/). 
      integer, intent(in) :: map_x_indices(:)
      integer, intent(out) :: ier

      ier = 0
      this % ndims = size(xmin)
      allocate(this % xmin(this % ndims))
      allocate(this % xmax(this % ndims))
      allocate(this % delta_x(this % ndims))
      allocate(this % M(this % ndims))
      allocate(this % map_x_indices(this % ndims))

      this % xmin = xmin
      this % xmax = xmax
      this % M = M
      this % map_x_indices = map_x_indices
      
      this % delta_x = (this%xmax - this%xmin)/this%M

    end subroutine xgrid_init

    subroutine xgrid_free(this, ier)

      ! Destructor

      implicit none
      type(xgrid_obj) :: this
      integer, intent(out) :: ier

      ier = 0
      deallocate(this % xmin)
      deallocate(this % xmax)
      deallocate(this % delta_x)
      deallocate(this % M)
      deallocate(this % map_x_indices)
      
    end subroutine xgrid_free

    subroutine xgrid_get(this, indices, x, ier)

      ! Given a set of indices, return the components of 
      ! x.

      implicit none
      type(xgrid_obj) :: this
      integer, intent(in) :: indices(:) ! from indexer
      real(r8), intent(out) :: x(:)     ! collocation position
      integer, intent(out) :: ier

      integer j

      ier = 0
      do j = 1, this % ndims
         x(j)  = this%xmin(j) + &
              &  indices(this%map_x_indices(j)) * this%delta_x(j)
      enddo
      
    end subroutine xgrid_get

    subroutine xgrid_test(this, ier)

      ! Tester

      implicit none
      type(xgrid_obj) :: this
      integer, intent(out) :: ier

      integer, parameter :: ndims = 2 ! 2-d space
      integer, parameter, dimension(ndims) :: M = (/5, 8/) ! no of collocation intervals
      ! computational domain
      real(r8), parameter, dimension(ndims) :: xmin = (/0._r8, -1._r8/)
      real(r8), parameter, dimension(ndims) :: xmax = (/1._r8, +1._r8/)
      integer :: map_x_indices(ndims), indices(ndims+1)
      real(r8) :: x(ndims)

      ier = 0

      print *,'testing xgrid ', trim(version)
      map_x_indices = (/2,3/)
      print *,'ndims=', ndims
      print *,'M=', M
      print '(2(a,2f10.3))','xmin=', xmin, ' xmax=', xmax
      print *,'index mapping=', map_x_indices 

      call xgrid_init(this, xmin, xmax, M, map_x_indices, ier)

      indices = (/1, 0, 0/)
      call xgrid_get(this, indices, x, ier)
      print *,'indices=', indices
      print '(a,2f10.3)', ' x=',x

      indices = (/2, 0, 0/)
      call xgrid_get(this, indices, x, ier)
      print *,'indices=', indices
      print '(a,2f10.3)',' x=',x
      
      indices = (/2, 5, 0/)
      call xgrid_get(this, indices, x, ier)
      print *,'indices=', indices 
      print '(a,2f10.3)',' x=',x

      indices = (/2, 0, 8/)
      call xgrid_get(this, indices, x, ier)
      print *,'indices=', indices
      print '(a,2f10.3)',' x=',x

      indices = (/2, 5, 8/)
      call xgrid_get(this, indices, x, ier)
      print *,'indices=', indices
      print '(a,2f10.3)',' x=',x

      call xgrid_free(this, ier)

    end subroutine xgrid_test

end module xgrid
