MPAS Init Atm Core for testing the MPAS KD Tree

mpas_init_atm_core.F

! Copyright (c) 2013,  Los Alamos National Security, LLC (LANS)
! and the University Corporation for Atmospheric Research (UCAR).
!
! Unless noted otherwise source code is licensed under the BSD license.
! Additional copyright and license information can be found in the LICENSE file
! distributed with this code, or at http://mpas-dev.github.com/license.html
!
module init_atm_core


   contains


   function init_atm_core_init(domain, startTimeStamp) result(ierr)
   
      use mpas_derived_types
      use mpas_stream_manager
      use mpas_io_streams, only : MPAS_STREAM_NEAREST
      use init_atm_cases
   
      implicit none
   
      type (domain_type), intent(inout) :: domain
      character(len=*), intent(out) :: startTimeStamp

      type (block_type), pointer :: block
      type (mpas_pool_type), pointer :: state, mesh 
      character (len=StrKIND), pointer :: xtime
      character (len=StrKIND), pointer :: initial_time
      character (len=StrKIND), pointer :: config_start_time
      real (kind=RKIND), pointer :: sphere_radius
      integer :: ierr


      ierr = 0

      block => domain % blocklist
      do while (associated(block))
         call mpas_pool_get_subpool(block % structs, 'state', state)
         call mpas_pool_get_subpool(block % structs, 'mesh', mesh)
         call mpas_pool_get_array(state, 'xtime', xtime)
         call mpas_pool_get_array(state, 'initial_time', initial_time)
         call mpas_pool_get_config(mesh, 'sphere_radius', sphere_radius)
         call mpas_pool_get_config(block % configs, 'config_start_time', config_start_time)

         startTimeStamp = config_start_time
         xtime = config_start_time
         initial_time = config_start_time
         domain % sphere_radius = a      ! Appears in output files
         sphere_radius = a               ! Used in setting up test cases

         block => block % next
      end do 

      call MPAS_stream_mgr_add_att(domain % streamManager, 'sphere_radius', domain % sphere_radius, streamID='output', ierr=ierr)
      call MPAS_stream_mgr_add_att(domain % streamManager, 'sphere_radius', domain % sphere_radius, streamID='surface', ierr=ierr)

      !
      ! We don't actually expect the time in the (most likely 'static') file to
      !    match the time in the namelist, so just read whatever time we find in
      !    the input file.
      !
      call MPAS_stream_mgr_read(domain % streamManager, whence=MPAS_STREAM_NEAREST, ierr=ierr)
      call MPAS_stream_mgr_reset_alarms(domain % streamManager, direction=MPAS_STREAM_INPUT, ierr=ierr)

   end function init_atm_core_init

   function init_atm_core_run(domain) result(ierr)
   
      use mpas_derived_types
      use mpas_stream_manager
      use mpas_timer
      use init_atm_cases
      use mpas_kd_tree
   
      implicit none
   
      type (domain_type), intent(inout) :: domain
      type (block_type), pointer :: block_ptr
      type (mpas_pool_type), pointer :: mesh

      real(kind=RKIND), dimension(:), pointer :: latCell
      real(kind=RKIND), dimension(:), pointer :: lonCell
      real(kind=RKIND), dimension(:), pointer :: xCell
      real(kind=RKIND), dimension(:), pointer :: yCell
      real(kind=RKIND), dimension(:), pointer :: zCell
      real(kind=RKIND), dimension(:,:), pointer :: array
      integer, pointer :: nCells

      type (mpas_kd_type), dimension(:), pointer :: points => null()
      type (mpas_kd_type), dimension(:), pointer :: points2 => null()
      type (mpas_kd_type), pointer :: tree => null()
      type (mpas_kd_type), pointer :: tree2 => null()
      type (mpas_kd_type), pointer :: res => null()
      type (mpas_kd_type), pointer :: res2 => null()

      real(kind=RKIND), dimension(3) :: min_point
      real(kind=RKIND) :: min_distance
      real(kind=RKIND) :: distance
      real(kind=RKIND) :: kd_distance

      real(kind=RKIND) :: dis1, dis2, mindist1, mindist2
      type (mpas_kd_type) :: minres1, minres2
      

      integer :: i, j
      integer :: num_missed
      integer :: ierr

      ierr = 0

      call mpas_log_write('====================================================================')
      call mpas_log_write('================== STARTING TEST OF MPAS_KD_TREE ====================')

      call mpas_pool_get_subpool(domain % blocklist % structs, 'mesh', mesh)
      call mpas_pool_get_dimension(domain % blocklist % dimensions, 'nCells', nCells)
      call mpas_pool_get_array(mesh, 'xCell', xCell)
      call mpas_pool_get_array(mesh, 'yCell', yCell)
      call mpas_pool_get_array(mesh, 'zCell', zCell)
      call mpas_pool_get_array(mesh, 'latCell', latCell)
      call mpas_pool_get_array(mesh, 'lonCell', lonCell)

      allocate(points(nCells))

      do i = 1, nCells
         allocate(points(i) % point(3))
         points(i) % point(:) = (/xCell(i), yCell(i), zCell(i)/)
         points(i) % id = i
      enddo
      
      ! Construct KD Tree and see if it was constructed succesfully
      tree => mpas_kd_construct(points, 3)
      if (.not. associated(tree)) then
         call mpas_log_write("ERROR: Tree was not associated when it should have been", messageType=MPAS_LOG_ERR)
      endif


      ! Search the KD-Tree for the same points
      do i = 1, nCells
         call mpas_kd_search(tree, (/xCell(i), yCell(i), zCell(i)/), res)
         if (.not. all((/xCell(i), yCell(i), zCell(i)/) == res % point(:))) then
            call mpas_log_write("ERROR: Result did not equal the query!", messageType=MPAS_LOG_ERR)
         endif
      enddo

      allocate(array(3, nCells))
      call random_number(array(:,:))
      array =  (array(:,:) * (1.0_RKIND + 1.0_RKIND - (-1.0_RKIND)) + (-1.0_RKIND))

      num_missed = 0

      do i = 1, nCells
         call mpas_kd_search(tree, real(array(:,i),kind=RKIND), res)

         ! Brute force to see if we get the right answer
         min_distance = huge(min_distance)
         do j = 1, nCells
            distance = sum(((/xCell(j), yCell(j), zCell(j)/) - real(array(:,i),kind=R8KIND))**2)
            if (distance < min_distance) then
               min_point = (/xCell(j), yCell(j), zCell(j)/)
               min_distance = distance
            endif
         enddo

         if (all(min_point(:) /= res % point(:))) then
            call mpas_log_write("ERROR: The point found via brute force was not the same as the one from the kd search", &
                                 messageType=MPAS_LOG_ERR)
            write(0,*) "min point: ", min_point(:) 
            write(0,*) "res: ", res % point(:)
            num_missed = num_missed + 1
         endif
      enddo

      call mpas_kd_free(tree)
      deallocate(points)
      deallocate(array)

    
      call mpas_log_write('Testing tie breaking...')

      !
      ! Single and double test - Create two points that are equal in single precision, but not
      ! in double.
      !
      allocate(points(4))
      allocate(points2(4))
      do i = 1, 4
         allocate(points(i) % point(2))
         allocate(points2(i) % point(2))
      enddo

      points(1) % point(:) = (/1.0, 1.0/)
      points(1) % id = 1
      points(2) % point(:) = (/2.0, 1.0/)
      points(2) % id = 2
      points(3) % point(:) = (/1.0, 2.0/) ! Halo Cell
      points(3) % id = -3
      points(4) % point(:) = (/2.0, 2.0/) ! Halo Cell
      points(4) % id = -4

      points2(1) % point(:) = (/1.0, 2.0/)
      points2(1) % id = 1
      points2(2) % point(:) = (/2.0, 2.0/)
      points2(2) % id = 2
      points2(3) % point(:) = (/1.0, 1.0/) ! Halo Cell
      points2(3) % id = -3
      points2(4) % point(:) = (/2.0, 1.0/) ! Halo Cell
      points2(4) % id = -4

      tree => mpas_kd_construct(points, 2)
      if (.not. associated(tree)) then
         call mpas_log_write("ERROR: Tree was not associated when it should have been", messageType=MPAS_LOG_ERR)
      endif

      tree2 => mpas_kd_construct(points2, 2)
      if (.not. associated(tree2)) then
         call mpas_log_write("ERROR: Tree was not associated when it should have been", messageType=MPAS_LOG_ERR)
      endif

      ! Find solution, the answer should be the point closets to the origin
      ! Is the above true for points on a sphere??
      mindist1 = huge(mindist1)
      mindist2 = huge(mindist2)
      do i = 1, 4
          ! Tree 1
          dis1 = sum((/0.0, 0.0/) - points(i) % point(:))**2
          if (dis1 < mindist1) then
              mindist1 = dis1
              minres1 = points(i)
          endif 

          ! Tree 2
          dis2 = sum((/0.0, 0.0/) - points2(i) % point(:)) **2
          if (dis2 < mindist2) then
              mindist2 = dis2
              minres2 = points2(i)
          endif
      enddo

      ! Search for the above
      call mpas_kd_search(tree, (/1.5, 1.5/), res)
      call mpas_kd_search(tree2, (/1.5, 1.5/), res2)
    
      if (all(res % point(:) /= minres1 % point(:))) then
         call mpas_log_write("ERROR: Tie was not borken correctly for the tree1", messageType=MPAS_LOG_ERR)
      endif

      if (all(res2 % point(:) /= minres1 % point(:))) then
         call mpas_log_write("ERROR: Tie was not borken correctly for the tree2", messageType=MPAS_LOG_ERR)
      endif

      call mpas_kd_free(tree)
      call mpas_kd_free(tree2)
      deallocate(points)
      deallocate(points2)

      call mpas_log_write('================== COMPLETED TEST OF MPAS_KD_TREE ====================')
   
   end function init_atm_core_run
   
   
   function init_atm_core_finalize(domain) result(ierr)
   
      use mpas_derived_types
      use mpas_decomp
      use mpas_stream_manager
      use mpas_log, only : mpas_log_write
   
      implicit none
   
      type (domain_type), intent(inout) :: domain 
      integer :: ierr


      ierr = 0

      call mpas_decomp_destroy_decomp_list(domain % decompositions)

      call mpas_log_write('')
      call mpas_log_write('********************************************************')
      call mpas_log_write('   Finished running the init_atmosphere core')
      call mpas_log_write('********************************************************')
   
   end function init_atm_core_finalize
   
end module init_atm_core