From a7e7de0d0b655294b0510ea89843dafcc71be729 Mon Sep 17 00:00:00 2001 From: Francis Vitt Date: Wed, 3 Sep 2025 15:46:00 -0600 Subject: [PATCH] untested code for lonlat2phys regrid --- src/utils/esmf_lonlat2phys_mod.F90 | 160 +++++++++++++++++++++++++++++ src/utils/esmf_phys2lonlat_mod.F90 | 12 +-- src/utils/remap.F90 | 8 ++ 3 files changed, 172 insertions(+), 8 deletions(-) create mode 100644 src/utils/esmf_lonlat2phys_mod.F90 diff --git a/src/utils/esmf_lonlat2phys_mod.F90 b/src/utils/esmf_lonlat2phys_mod.F90 new file mode 100644 index 0000000000..63a388d10a --- /dev/null +++ b/src/utils/esmf_lonlat2phys_mod.F90 @@ -0,0 +1,160 @@ +!------------------------------------------------------------------------------ +! Provides methods for mapping from regular longitude / latitude grid +! to physics grid to via ESMF regridding capabilities +!------------------------------------------------------------------------------ +module esmf_lonlat2phys_mod + use shr_kind_mod, only: r8 => shr_kind_r8 + use cam_logfile, only: iulog + use cam_abortutils, only: endrun + use spmd_utils, only: masterproc + use ppgrid, only: pver + + use ESMF, only: ESMF_RouteHandle, ESMF_Field, ESMF_ArraySpec, ESMF_ArraySpecSet + use ESMF, only: ESMF_FieldCreate, ESMF_FieldRegridStore + use ESMF, only: ESMF_FieldGet, ESMF_FieldRegrid + use ESMF, only: ESMF_KIND_I4, ESMF_KIND_R8, ESMF_TYPEKIND_R8 + use ESMF, only: ESMF_REGRIDMETHOD_BILINEAR, ESMF_POLEMETHOD_NONE, ESMF_EXTRAPMETHOD_NEAREST_IDAVG + use ESMF, only: ESMF_TERMORDER_SRCSEQ, ESMF_MESHLOC_ELEMENT, ESMF_STAGGERLOC_CENTER + use ESMF, only: ESMF_FieldDestroy, ESMF_RouteHandleDestroy + use esmf_check_error_mod, only: check_esmf_error + + implicit none + + private + + public :: esmf_lonlat2phys_init + public :: esmf_lonlat2phys_regrid + public :: esmf_lonlat2phys_destroy + public :: fields_bundle_t + public :: n_flx_flds + + type(ESMF_RouteHandle) :: rh_lonlat2phys_3d + + type(ESMF_Field) :: physfld_3d + type(ESMF_Field) :: lonlatfld_3d + + interface esmf_lonlat2phys_regrid + module procedure esmf_lonlat2phys_regrid_3d + end interface esmf_lonlat2phys_regrid + + type :: fields_bundle_t + real(r8), pointer :: fld(:,:,:) => null() + end type fields_bundle_t + + integer, parameter :: n_flx_flds = 2 ! 3D uflux and vflux + +contains + + !------------------------------------------------------------------------------ + !------------------------------------------------------------------------------ + subroutine esmf_lonlat2phys_init() + use esmf_phys_mesh_mod, only: physics_grid_mesh + use esmf_lonlat_grid_mod, only: lonlat_grid + + type(ESMF_ArraySpec) :: arrayspec + integer(ESMF_KIND_I4), pointer :: factorIndexList(:,:) + real(ESMF_KIND_R8), pointer :: factorList(:) + integer :: smm_srctermproc, smm_pipelinedep, rc + + character(len=*), parameter :: subname = 'esmf_lonlat2phys_init: ' + + smm_srctermproc = 0 + smm_pipelinedep = 16 + + ! create ESMF fields + + ! 3D phys fld + call ESMF_ArraySpecSet(arrayspec, 3, ESMF_TYPEKIND_R8, rc=rc) + call check_esmf_error(rc, subname//'ESMF_ArraySpecSet 3D phys fld ERROR') + + physfld_3d = ESMF_FieldCreate(physics_grid_mesh, arrayspec, & + gridToFieldMap=(/3/), meshloc=ESMF_MESHLOC_ELEMENT, & + ungriddedLBound=(/1,1/), ungriddedUBound=(/pver,n_flx_flds/), rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldCreate 3D phys fld ERROR') + + ! 3D lon lat grid + call ESMF_ArraySpecSet(arrayspec, 4, ESMF_TYPEKIND_R8, rc=rc) + call check_esmf_error(rc, subname//'ESMF_ArraySpecSet 3D lonlat fld ERROR') + + lonlatfld_3d = ESMF_FieldCreate( lonlat_grid, arrayspec, staggerloc=ESMF_STAGGERLOC_CENTER, & + ungriddedLBound=(/1,1/), ungriddedUBound=(/pver,n_flx_flds/), rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldCreate 3D lonlat fld ERROR') + + call ESMF_FieldRegridStore(srcField=lonlatfld_3d, dstField=physfld_3d, & + regridMethod=ESMF_REGRIDMETHOD_BILINEAR, & + polemethod=ESMF_POLEMETHOD_NONE, & + extrapMethod=ESMF_EXTRAPMETHOD_NEAREST_IDAVG, & + routeHandle=rh_lonlat2phys_3d, factorIndexList=factorIndexList, & + factorList=factorList, srcTermProcessing=smm_srctermproc, & + pipelineDepth=smm_pipelinedep, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldRegridStore 3D routehandle ERROR') + + end subroutine esmf_lonlat2phys_init + + !------------------------------------------------------------------------------ + !------------------------------------------------------------------------------ + subroutine esmf_lonlat2phys_regrid_3d(lonlatflds, physflds) + use esmf_lonlat_grid_mod, only: lon_beg,lon_end,lat_beg,lat_end + use ppgrid, only: pcols, pver, begchunk, endchunk + use phys_grid, only: get_ncols_p + + type(fields_bundle_t), intent(in) :: lonlatflds(n_flx_flds) + type(fields_bundle_t), intent(inout) :: physflds(n_flx_flds) + + integer :: i, ichnk, ncol, ifld, ilev, icol, rc + real(ESMF_KIND_R8), pointer :: physptr(:,:,:) + real(ESMF_KIND_R8), pointer :: lonlatptr(:,:,:,:) + + character(len=*), parameter :: subname = 'esmf_lonlat2phys_regrid_3d: ' + + ! set values of lonlatfld_3d ESMF field + call ESMF_FieldGet(lonlatfld_3d, localDe=0, farrayPtr=lonlatptr, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldGet lonlatptr') + + do ifld = 1,n_flx_flds + lonlatptr(lon_beg:lon_end,lat_beg:lat_end,1:pver,ifld) = lonlatflds(ifld)%fld(lon_beg:lon_end,lat_beg:lat_end,1:pver) + end do + + ! regrid + call ESMF_FieldRegrid(lonlatfld_3d, physfld_3d, rh_lonlat2phys_3d, & + termorderflag=ESMF_TERMORDER_SRCSEQ, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldRegrid lonlatfld_3d->physfld_3d') + + ! get values from physfld_3d ESMF field + call ESMF_FieldGet(physfld_3d, localDe=0, farrayPtr=physptr, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldGet physptr') + + i = 0 + do ichnk = begchunk, endchunk + ncol = get_ncols_p(ichnk) + do icol = 1,ncol + i = i+1 + do ifld = 1,n_flx_flds + do ilev = 1,pver + physflds(ifld)%fld(ilev,icol,ichnk) = physptr(ilev,ifld,i) + end do + end do + end do + end do + + end subroutine esmf_lonlat2phys_regrid_3d + + !------------------------------------------------------------------------------ + !------------------------------------------------------------------------------ + subroutine esmf_lonlat2phys_destroy() + + integer :: rc + character(len=*), parameter :: subname = 'esmf_lonlat2phys_destroy: ' + + call ESMF_RouteHandleDestroy(rh_lonlat2phys_3d, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldDestroy rh_lonlat2phys_3d') + + call ESMF_FieldDestroy(lonlatfld_3d, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldDestroy lonlatfld_3d') + + call ESMF_FieldDestroy(physfld_3d, rc=rc) + call check_esmf_error(rc, subname//'ESMF_FieldDestroy physfld_3d') + + end subroutine esmf_lonlat2phys_destroy + +end module esmf_lonlat2phys_mod diff --git a/src/utils/esmf_phys2lonlat_mod.F90 b/src/utils/esmf_phys2lonlat_mod.F90 index a191cf31e5..bf18e3b6ef 100644 --- a/src/utils/esmf_phys2lonlat_mod.F90 +++ b/src/utils/esmf_phys2lonlat_mod.F90 @@ -33,11 +33,9 @@ module esmf_phys2lonlat_mod type(ESMF_Field) :: physfld_3d type(ESMF_Field) :: lonlatfld_3d - type(ESMF_Field) :: lonlatfld_3d_copy type(ESMF_Field) :: physfld_2d type(ESMF_Field) :: lonlatfld_2d - type(ESMF_Field) :: lonlatfld_2d_copy interface esmf_phys2lonlat_regrid module procedure esmf_phys2lonlat_regrid_2d @@ -86,7 +84,6 @@ subroutine esmf_phys2lonlat_init() lonlatfld_3d = ESMF_FieldCreate( lonlat_grid, arrayspec, staggerloc=ESMF_STAGGERLOC_CENTER, & ungriddedLBound=(/1,1/), ungriddedUBound=(/pver,nflds/), rc=rc) call check_esmf_error(rc, subname//'ESMF_FieldCreate 3D lonlat fld ERROR') - lonlatfld_3d_copy = lonlatfld_3d ! 2D phys fld call ESMF_ArraySpecSet(arrayspec, 1, ESMF_TYPEKIND_R8, rc=rc) @@ -102,20 +99,19 @@ subroutine esmf_phys2lonlat_init() lonlatfld_2d = ESMF_FieldCreate( lonlat_grid, arrayspec, staggerloc=ESMF_STAGGERLOC_CENTER, rc=rc) call check_esmf_error(rc, subname//'ESMF_FieldCreate 2D lonlat fld ERROR') - lonlatfld_2d_copy = lonlatfld_2d - call ESMF_FieldRegridStore(srcField=physfld_3d, dstField=lonlatfld_3d_copy, & + call ESMF_FieldRegridStore(srcField=physfld_3d, dstField=lonlatfld_3d, & regridMethod=ESMF_REGRIDMETHOD_BILINEAR, & - polemethod=ESMF_POLEMETHOD_NONE, & + polemethod=ESMF_POLEMETHOD_NONE, & extrapMethod=ESMF_EXTRAPMETHOD_NEAREST_IDAVG, & routeHandle=rh_phys2lonlat_3d, factorIndexList=factorIndexList, & factorList=factorList, srcTermProcessing=smm_srctermproc, & pipelineDepth=smm_pipelinedep, rc=rc) call check_esmf_error(rc, subname//'ESMF_FieldRegridStore 3D routehandle ERROR') - call ESMF_FieldRegridStore(srcField=physfld_2d, dstField=lonlatfld_2d_copy, & + call ESMF_FieldRegridStore(srcField=physfld_2d, dstField=lonlatfld_2d, & regridMethod=ESMF_REGRIDMETHOD_BILINEAR, & - polemethod=ESMF_POLEMETHOD_NONE, & + polemethod=ESMF_POLEMETHOD_NONE, & extrapMethod=ESMF_EXTRAPMETHOD_NEAREST_IDAVG, & routeHandle=rh_phys2lonlat_2d, factorIndexList=factorIndexList, & factorList=factorList, srcTermProcessing=smm_srctermproc, & diff --git a/src/utils/remap.F90 b/src/utils/remap.F90 index 7c2ae82c8c..af9ab9529f 100644 --- a/src/utils/remap.F90 +++ b/src/utils/remap.F90 @@ -34,6 +34,7 @@ subroutine nlgw_regrid_init() use esmf_lonlat_grid_mod, only: esmf_lonlat_grid_init use esmf_phys_mesh_mod, only: esmf_phys_mesh_init use esmf_phys2lonlat_mod, only: esmf_phys2lonlat_init + use esmf_lonlat2phys_mod, only: esmf_lonlat2phys_init integer, parameter :: reg_decomp = 332 @@ -50,6 +51,7 @@ subroutine nlgw_regrid_init() call esmf_lonlat_grid_init(64, 128) call esmf_phys_mesh_init() call esmf_phys2lonlat_init() + call esmf_lonlat2phys_init() ! for the lon-lat grid allocate(grid_map(4, ((endlon - beglon + 1) * (endlat - beglat + 1))), stat=astat) @@ -116,6 +118,7 @@ subroutine nlgw_regrid(phys_state) use esmf_zonal_mean_mod, only: esmf_zonal_mean_calc, esmf_zonal_mean_wsums, esmf_zonal_mean_masked use interpolate_data, only: lininterp use esmf_phys2lonlat_mod, only: fields_bundle_t, nflds + use esmf_lonlat2phys_mod, only: esmf_lonlat2phys_regrid, n_flx_flds use mpishorthand type(physics_state), intent(in) :: phys_state(begchunk:endchunk) @@ -153,6 +156,9 @@ subroutine nlgw_regrid(phys_state) type(fields_bundle_t) :: physflds(nflds) type(fields_bundle_t) :: lonlatflds(nflds) + type(fields_bundle_t) :: phys_flx_flds(n_flx_flds) + type(fields_bundle_t) :: lonlat_flx_flds(n_flx_flds) + call t_startf('nlgw_gather') call t_startf('nlgw_unchunk') @@ -269,10 +275,12 @@ end subroutine nlgw_regrid !----------------------------------------------------------------------------- subroutine nlgw_regrid_final() use esmf_phys2lonlat_mod, only: esmf_phys2lonlat_destroy + use esmf_lonlat2phys_mod, only: esmf_lonlat2phys_destroy use esmf_lonlat_grid_mod, only: esmf_lonlat_grid_destroy use esmf_phys_mesh_mod, only: esmf_phys_mesh_destroy call esmf_phys2lonlat_destroy() + call esmf_lonlat2phys_destroy() call esmf_lonlat_grid_destroy() call esmf_phys_mesh_destroy()