test_regrid2xg Subroutine

subroutine test_regrid2xg(atm_nx, atm_ny, ocn_nx, ocn_ny, atm_dx, atm_dy, ocn_dx, ocn_dy, rc)

Arguments

Type IntentOptional Attributes Name
integer, intent(in) :: atm_nx
integer, intent(in) :: atm_ny
integer, intent(in) :: ocn_nx
integer, intent(in) :: ocn_ny
real(kind=ESMF_KIND_R4), intent(in) :: atm_dx
real(kind=ESMF_KIND_R4), intent(in) :: atm_dy
real(kind=ESMF_KIND_R4), intent(in) :: ocn_dx
real(kind=ESMF_KIND_R4), intent(in) :: ocn_dy
integer, intent(out) :: rc

Source Code

  subroutine test_regrid2xg(atm_nx, atm_ny, ocn_nx, ocn_ny, atm_dx, atm_dy, &
    ocn_dx, ocn_dy, rc)
    ! arguments
    integer, intent(in)             :: atm_nx, atm_ny, ocn_nx, ocn_ny
    real(ESMF_KIND_R4), intent(in)  :: atm_dx, atm_dy, ocn_dx, ocn_dy
    integer, intent(out)            :: rc

    ! local variables
    type(ESMF_Grid)                 :: grid_atm, grid_ocn
    type(ESMF_Field)                :: f_atm, f_ocn, f_xgrid
    real(ESMF_KIND_R8)              :: startx, starty
    integer                         :: localrc, npet, i, j, lpet
    real(ESMF_KIND_R8), pointer     :: weights(:)
    integer(ESMF_KIND_I4), pointer  :: indices(:,:)
    real(ESMF_KIND_R8), pointer     :: coordX(:), coordY(:)
    type(ESMF_XGrid)                :: xgrid
    type(ESMF_XGridSpec)            :: sparseMatA2X(1)
    integer                         :: gn(2), simax, dimax
    real(ESMF_KIND_R8), pointer     :: atm(:,:), ocn(:,:), exf(:)
    type(ESMF_RouteHandle)          :: rh
    type(ESMF_DistGrid)             :: distgridM

    type(ESMF_Field)                :: fa_atm, fa_ocn, fa_xgrid
    type(ESMF_Field)                :: aa_atm, aa_ocn, aa_xgrid
    type(ESMF_Mesh)                 :: mesh_atm, mesh_ocn, mesh_xgrid
    
    type(ESMF_VM)   :: vm

    localrc = ESMF_SUCCESS

    call ESMF_VMGetCurrent(vm, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_VMGet(vm, petCount=npet, localPet=lpet, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! Grid constants
    ! These are input arguments from upper level now
    !atm_nx = 10
    !atm_ny = 10
    !atm_dx = 0.1
    !atm_dy = 0.1
    !ocn_nx = 14   ! change to 100 to force a match
    !ocn_ny = 14   ! change to 100 to force a match
    !ocn_dx = 0.06
    !ocn_dy = 0.06
    
    !------------- ATM --------------
    ! atm grid, horizontally decomposed
    grid_atm = ESMF_GridCreateNoPeriDim(maxIndex=(/atm_nx, atm_ny/), &
      indexflag=ESMF_INDEX_GLOBAL, &
      !gridEdgeLWidth=(/0,0/), gridEdgeUWidth=(/0,0/), &
      regDecomp=(/npet, 1/), &
      coordDep1=(/1/), &
      coordDep2=(/2/), &
      rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_GridAddCoord(grid_atm, staggerloc=ESMF_STAGGERLOC_CENTER, &
        rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    call ESMF_GridAddCoord(grid_atm, staggerloc=ESMF_STAGGERLOC_CORNER, &
        rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! global indexing
    ! atm grid is not decomposed in the y direction
    !startx = lpet*atm_nx/npet*atm_dx
    startx = 0.
    starty = 0.
    ! compute coord
    ! X center
    call ESMF_GridGetCoord(grid_atm, localDE=0, staggerLoc=ESMF_STAGGERLOC_CENTER, &
        coordDim=1, farrayPtr=coordX, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordX,1), ubound(coordX,1)
      coordX(i) = startx + atm_dx/2. + (i-1)*atm_dx
    enddo
    ! X corner
    call ESMF_GridGetCoord(grid_atm, localDE=0, staggerLoc=ESMF_STAGGERLOC_CORNER, &
        coordDim=1, farrayPtr=coordX, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordX,1), ubound(coordX,1)
      coordX(i) = startx + (i-1)*atm_dx
    enddo
    !print *, 'startx: ', startx, lbound(coordX, 1), 'coordX: ', coordX
    ! Y center
    call ESMF_GridGetCoord(grid_atm, localDE=0, staggerLoc=ESMF_STAGGERLOC_CENTER, &
        coordDim=2, farrayPtr=coordY, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordY,1), ubound(coordY,1)
      coordY(i) = starty + atm_dy/2. + (i-1)*atm_dy
    enddo
    ! Y corner
    call ESMF_GridGetCoord(grid_atm, localDE=0, staggerLoc=ESMF_STAGGERLOC_CORNER, &
        coordDim=2, farrayPtr=coordY, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordY,1), ubound(coordY,1)
      coordY(i) = starty + (i-1)*atm_dy
    enddo

    !------------- OCN --------------
    ! ocn grid, horizontally decomposed
    grid_ocn = ESMF_GridCreateNoPeriDim(maxIndex=(/ocn_nx, ocn_ny/), &
      indexflag=ESMF_INDEX_GLOBAL, &
      !gridEdgeLWidth=(/0,0/), gridEdgeUWidth=(/0,0/), &
      regDecomp=(/npet, 1/), &
      coordDep1=(/1/), &
      coordDep2=(/2/), &
      rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_GridAddCoord(grid_ocn, staggerloc=ESMF_STAGGERLOC_CENTER, &
        rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    call ESMF_GridAddCoord(grid_ocn, staggerloc=ESMF_STAGGERLOC_CORNER, &
        rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! ocn grid is not decomposed in the y direction
    !startx = lpet*ocn_nx/npet*ocn_dx
    startx = 0.
    starty = 0.
    ! compute coord
    ! X center
    call ESMF_GridGetCoord(grid_ocn, localDE=0, staggerLoc=ESMF_STAGGERLOC_CENTER, &
        coordDim=1, farrayPtr=coordX, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordX,1), ubound(coordX,1)
      coordX(i) = startx + ocn_dx/2. + (i-1)*ocn_dx
    enddo
    ! X corner
    call ESMF_GridGetCoord(grid_ocn, localDE=0, staggerLoc=ESMF_STAGGERLOC_CORNER, &
        coordDim=1, farrayPtr=coordX, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordX,1), ubound(coordX,1)
      coordX(i) = startx + (i-1)*ocn_dx
    enddo
    ! Y center 
    call ESMF_GridGetCoord(grid_ocn, localDE=0, staggerLoc=ESMF_STAGGERLOC_CENTER, &
        coordDim=2, farrayPtr=coordY, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordY,1), ubound(coordY,1)
      coordY(i) = starty + ocn_dy/2. + (i-1)*ocn_dy
    enddo
    ! Y corner
    call ESMF_GridGetCoord(grid_ocn, localDE=0, staggerLoc=ESMF_STAGGERLOC_CORNER, &
        coordDim=2, farrayPtr=coordY, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    do i = lbound(coordY,1), ubound(coordY,1)
      coordY(i) = starty + (i-1)*ocn_dy
    enddo

    ! build Fields on the Grids
    f_atm = ESMF_FieldCreate(grid_atm, typekind=ESMF_TYPEKIND_R8, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    f_ocn = ESMF_FieldCreate(grid_ocn, typekind=ESMF_TYPEKIND_R8, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! Call RegridStore here to compute SMM weights and indices
    call ESMF_FieldRegridStore(srcField=f_atm, dstField=f_ocn, &
      regridmethod=ESMF_REGRIDMETHOD_CONSERVE, &
      unmappedaction = ESMF_UNMAPPEDACTION_IGNORE, &
      factorIndexList=indices, factorList=weights, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! make sure the numbers are consistent
    print *, lpet, size(weights), '-', weights
    print *, lpet, size(indices,1),'-', size(indices,2)
    do j = 1, size(indices,2)
         print *, indices(1,j), '->', indices(2,j)
    enddo

    simax = maxval(indices(1,:))
    dimax = maxval(indices(2,:))
    call ESMF_VMAllReduce(vm, (/ simax, dimax /), gn, 2, ESMF_REDUCE_MAX, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    print *, gn(1), gn(2)
    if(gn(2) /= ocn_nx*ocn_ny) then
      call ESMF_LogSetError(ESMF_RC_VAL_WRONG, ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)
      return
    endif

    ! analyze the returned the weights based on area information

    ! use the weights to generate an XGrid
    allocate(sparseMatA2X(1)%factorIndexList(size(indices,1), size(indices,2)))
    do i = 1, size(indices,1)
      do j = 1, size(indices,2)
         sparseMatA2X(1)%factorIndexList(i,j) = indices(i,j)
      enddo
    enddo
    sparseMatA2X(1)%factorList=>weights
    xgrid = ESMF_XGridCreateFromSparseMat(sideAGrid=(/grid_atm/), sideBGrid=(/grid_ocn/), &
        sparseMatA2X=sparseMatA2X, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return


    call ESMF_XGridGet(xgrid, distgridM=distgridM, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    call ESMF_DistGridPrint(distgridM, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! create a Field on the xgrid
    f_xgrid = ESMF_FieldCreate(xgrid=xgrid, TYPEKIND=ESMF_TYPEKIND_R8, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! regrid through the xgrid
    ! set up src flux
    call ESMF_FieldGet(f_atm, farrayPtr=atm, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    call ESMF_FieldGet(f_xgrid, farrayPtr=exf, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! compute regrid routehandle
    call ESMF_FieldRegridStore(xgrid, f_atm, f_xgrid, routehandle=rh, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    ! execute regrid
    ! once area information is retrieved from mesh regrid store, then
    ! one can perform variable flux exchange and check flux conservation
    ! for now the src flux is a constant term, this results in const dst flux
    atm = 2.
    call ESMF_FieldRegrid(f_atm, f_xgrid, routehandle=rh, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    print *, lpet, exf

    ! clean up
    call ESMF_FieldDestroy(f_atm, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_FieldDestroy(f_ocn, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_GridDestroy(grid_atm, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_GridDestroy(grid_ocn, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    call ESMF_FieldDestroy(f_xgrid, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return
    call ESMF_XGridDestroy(xgrid, rc=localrc)
    if (ESMF_LogFoundError(localrc, &
        ESMF_ERR_PASSTHRU, &
        ESMF_CONTEXT, rcToReturn=rc)) return

    deallocate(sparseMatA2X(1)%factorIndexList)

    rc = ESMF_SUCCESS

  end subroutine test_regrid2xg