@@ -23,6 +23,7 @@ module lnd_comp_domain
2323 use ESMF , only : ESMF_RouteHandleDestroy, ESMF_GridGet, ESMF_GridGetCoord
2424 use ESMF , only : ESMF_FieldRegridGetArea, ESMF_CoordSys_Flag
2525 use ESMF , only : ESMF_MeshGetFieldBounds, ESMF_COORDSYS_CART, ESMF_KIND_R8
26+ use ESMF , only : ESMF_MeshDestroy, ESMF_DistGridCreate, ESMF_VMAllGatherV
2627 use NUOPC, only : NUOPC_CompAttributeGet
2728
2829 use lnd_comp_kind , only : r4 = > shr_kind_r4
@@ -62,10 +63,11 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
6263 ! local variables
6364 real (r4 ), target , allocatable :: tmpr4 (:)
6465 real (r8 ), target , allocatable :: tmpr8 (:)
65- integer :: n
66+ integer :: n, petCount, localPet
6667 integer :: decomptile(2 ,6 )
6768 integer :: maxIndex(2 )
6869 type (ESMF_Decomp_Flag) :: decompflagPTile(2 ,6 )
70+ type (ESMF_VM) :: vm
6971
7072 type (field_type) :: flds(1 )
7173 integer :: numOwnedElements, spatialDim, rank
@@ -87,6 +89,16 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
8789 rc = ESMF_SUCCESS
8890 call ESMF_LogWrite(subname// ' called'
8991
92+ ! ---------------------
93+ ! Query VM
94+ ! ---------------------
95+
96+ call ESMF_GridCompGet(gcomp, vm= vm, rc= rc)
97+ if (ChkErr(rc,__LINE__,u_FILE_u)) return
98+
99+ call ESMF_VMGet(vm, petCount= petCount, localPet= localPet, rc= rc)
100+ if (ChkErr(rc,__LINE__,u_FILE_u)) return
101+
90102 ! ---------------------
91103 ! Set decomposition and decide it is regional or global
92104 ! ---------------------
@@ -101,7 +113,14 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
101113 ! set number of tiles
102114 noahmp% domain% ntiles = 6
103115
104- ! set decomposition
116+ ! check user provided layout
117+ if (petCount /= noahmp% domain% layout(1 )* noahmp% domain% layout(2 )* noahmp% domain% ntiles) then
118+ call ESMF_LogWrite(trim (subname)// " : ERROR in layout. layout_x * layout_y * 6 != #PETs"
119+ rc = ESMF_FAILURE
120+ return
121+ end if
122+
123+ ! use user provided layout to set decomposition
105124 do n = 1 , noahmp% domain% ntiles
106125 decomptile(1 ,n) = noahmp% domain% layout(1 )
107126 decomptile(2 ,n) = noahmp% domain% layout(2 )
@@ -226,7 +245,7 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
226245 if (ChkErr(rc,__LINE__,u_FILE_u)) return
227246 vegtype(:) = int (tmpr8 )
228247 else
229- write (filename, fmt= " (A,I0,A)" trim (noahmp% nmlist% input_dir )// ' C' % domain% ni, ' .vegetation_type.tile*.nc'
248+ write (filename, fmt= " (A,I0,A)" trim (noahmp% nmlist% fixed_dir )// ' C' % domain% ni, ' .vegetation_type.tile*.nc'
230249 flds(1 )% short_name = ' vegetation_type'
231250 flds(1 )% ptr1r4 = > tmpr4
232251 call read_tiled_file(noahmp, filename, flds, rc= rc)
@@ -257,6 +276,16 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
257276 call ESMF_MeshSet(noahmp% domain% mesh, elementMask= noahmp% domain% mask, rc= rc)
258277 if (chkerr(rc,__LINE__,u_FILE_u)) return
259278
279+ ! ---------------------
280+ ! Modify decomposition to evenly distribute land and ocean points
281+ ! ---------------------
282+
283+ if (trim (noahmp% nmlist% decomp_type) == ' custom' then
284+ ! modify decomposition
285+ call lnd_modify_decomp(gcomp, noahmp, rc)
286+ if (chkerr(rc,__LINE__,u_FILE_u)) return
287+ end if
288+
260289 ! ---------------------
261290 ! Get height from orography file
262291 ! ---------------------
@@ -339,4 +368,221 @@ subroutine lnd_set_decomp_and_domain_from_mosaic(gcomp, noahmp, rc)
339368
340369 end subroutine lnd_set_decomp_and_domain_from_mosaic
341370
371+ ! ===============================================================================
372+ subroutine lnd_modify_decomp (gcomp , noahmp , rc )
373+
374+ ! input/output variables
375+ type (ESMF_GridComp), intent (in ) :: gcomp
376+ type (noahmp_type), intent (inout ) :: noahmp
377+ integer , intent (out ) :: rc
378+
379+ ! local variables
380+ type (ESMF_VM) :: vm
381+ type (ESMF_Mesh) :: mesh
382+ type (ESMF_DistGrid) :: distgrid, distgrid_new
383+ type (field_type) :: flds(1 )
384+ integer :: n, m, g
385+ integer :: petCount, localPet
386+ integer :: lsize, gsize(1 )
387+ integer :: nlnd, nocn
388+ integer :: begl_l, endl_l, begl_o, endl_o
389+ integer , allocatable :: nlnd_loc(:)
390+ integer , allocatable :: nocn_loc(:)
391+ integer , allocatable :: mask_glb(:)
392+ integer , allocatable :: gindex_loc(:)
393+ integer , allocatable :: gindex_glb(:)
394+ integer , allocatable :: gindex_lnd(:)
395+ integer , allocatable :: gindex_ocn(:)
396+ integer , allocatable :: gindex_new(:)
397+ integer , allocatable :: lsize_arr(:)
398+ real (r4 ), target , allocatable :: tmpr4 (:)
399+ character (len= cl) :: msg, filename
400+ character (len=* ), parameter :: subname = trim (modName)// ' :(lnd_modify_decomp) '
401+ !- ------------------------------------------------------------------------------
402+
403+ rc = ESMF_SUCCESS
404+ call ESMF_LogWrite(subname// ' called'
405+
406+ ! ---------------------
407+ ! Query VM
408+ ! ---------------------
409+
410+ call ESMF_GridCompGet(gcomp, vm= vm, rc= rc)
411+ if (ChkErr(rc,__LINE__,u_FILE_u)) return
412+
413+ call ESMF_VMGet(vm, petCount= petCount, localPet= localPet, rc= rc)
414+ if (ChkErr(rc,__LINE__,u_FILE_u)) return
415+
416+ ! ---------------------
417+ ! Query existing mesh
418+ ! ---------------------
419+
420+ ! retrive default distgrid
421+ call ESMF_MeshGet(noahmp% domain% mesh, elementdistGrid= distgrid, rc= rc)
422+ if (chkerr(rc,__LINE__,u_FILE_u)) return
423+
424+ ! get local number of elements
425+ call ESMF_DistGridGet(distgrid, localDe= 0 , elementCount= lsize, rc= rc)
426+ if (chkerr(rc,__LINE__,u_FILE_u)) return
427+
428+ ! calculate number of elements globally
429+ call ESMF_VMAllReduce(vm, (/ lsize / ), gsize, 1 , ESMF_REDUCE_SUM, rc= rc)
430+ if (chkerr(rc,__LINE__,u_FILE_u)) return
431+
432+ nlnd = count (noahmp% domain% mask(:) > 0 , dim= 1 )
433+ nocn = lsize- nlnd
434+ write (msg, fmt= ' (A,4I8)' trim (subname)// ' : lsize, gsize, nlnd, nocn = '
435+ lsize, gsize(1 ), nlnd, nocn
436+ call ESMF_LogWrite(trim (msg), ESMF_LOGMSG_INFO)
437+
438+ ! get default sequence index
439+ allocate (gindex_loc(lsize))
440+ call ESMF_DistGridGet(distgrid, 0 , seqIndexList= gindex_loc, rc= rc)
441+ if (chkerr(rc,__LINE__,u_FILE_u)) return
442+
443+ ! ---------------------
444+ ! Create global view of indexes and mask
445+ ! --------------------
446+
447+ ! collect local sizes
448+ allocate (lsize_arr(petCount))
449+ call ESMF_VMAllGatherV(vm, sendData= (/ lsize / ), sendCount= 1 , &
450+ recvData= lsize_arr, recvCounts= (/ (1 , n = 0 , petCount-1 ) / ), &
451+ recvOffsets= (/ (n, n = 0 , petCount-1 ) / ), rc= rc)
452+ if (chkerr(rc,__LINE__,u_FILE_u)) return
453+
454+ ! global view of indexes
455+ allocate (gindex_glb(gsize(1 )))
456+ gindex_glb(:) = 0
457+ call ESMF_VMAllGatherV(vm, sendData= gindex_loc, sendCount= lsize, &
458+ recvData= gindex_glb, recvCounts= lsize_arr, &
459+ recvOffsets= (/ (sum (lsize_arr(1 :n))- lsize_arr(1 ), n = 1 , petCount) / ), rc= rc)
460+ if (chkerr(rc,__LINE__,u_FILE_u)) return
461+
462+ ! global view of land sea mask
463+ allocate (mask_glb(gsize(1 )))
464+ mask_glb(:) = 0
465+ call ESMF_VMAllGatherV(vm, sendData= noahmp% domain% mask, sendCount= lsize, &
466+ recvData= mask_glb, recvCounts= lsize_arr, &
467+ recvOffsets= (/ (sum (lsize_arr(1 :n))- lsize_arr(1 ), n = 1 , petCount) / ), rc= rc)
468+ if (chkerr(rc,__LINE__,u_FILE_u)) return
469+
470+ ! ---------------------
471+ ! split global indexes as land and ocean
472+ ! ---------------------
473+
474+ nlnd = count (mask_glb > 0 , dim= 1 )
475+ nocn = gsize(1 )- nlnd
476+ allocate (gindex_lnd(nlnd))
477+ gindex_lnd = 0
478+ allocate (gindex_ocn(nocn))
479+ gindex_ocn = 0
480+
481+ n = 0
482+ m = 0
483+ do g = 1 , gsize(1 )
484+ if (mask_glb(g) > 0 ) then
485+ n = n+1
486+ gindex_lnd(n) = gindex_glb(g)
487+ else
488+ m = m+1
489+ gindex_ocn(m) = gindex_glb(g)
490+ end if
491+ end do
492+
493+ ! ---------------------
494+ ! create new local indexes
495+ ! ---------------------
496+
497+ allocate (nlnd_loc(0 :petCount-1 ))
498+ nlnd_loc = 0
499+ allocate (nocn_loc(0 :petCount-1 ))
500+ nocn_loc = 0
501+ do n = 0 , petCount-1
502+ nlnd_loc(n) = nlnd/ petCount
503+ nocn_loc(n) = nocn/ petCount
504+ if (n < mod (nlnd, petCount)) then
505+ nlnd_loc(n) = nlnd_loc(n)+ 1
506+ else
507+ nocn_loc(n) = nocn_loc(n)+ 1
508+ end if
509+ end do
510+ if (localPet == 0 ) then
511+ begl_l = 1
512+ begl_o = 1
513+ else
514+ begl_l = sum (nlnd_loc(0 :localPet-1 ))+ 1
515+ begl_o = sum (nocn_loc(0 :localPet-1 ))+ 1
516+ end if
517+ endl_l = sum (nlnd_loc(0 :localPet))
518+ endl_o = sum (nocn_loc(0 :localPet))
519+
520+ write (msg,' (A,6I8)' trim (subname)// ' : nlnd_loc, nocn_loc, begl_l, endl_l, begl_o, endl_o = '
521+ nlnd_loc(localPet), nocn_loc(localPet), begl_l, endl_l, begl_o, endl_o
522+ call ESMF_LogWrite(trim (msg), ESMF_LOGMSG_INFO)
523+
524+ allocate (gindex_new(nlnd_loc(localPet)+ nocn_loc(localPet)))
525+ gindex_new(:nlnd_loc(localPet)) = gindex_lnd(begl_l:begl_l)
526+ gindex_new(nlnd_loc(localPet)+ 1 :) = gindex_ocn(begl_o:endl_o)
527+
528+ if (noahmp% nmlist% debug_level > 10 ) then
529+ do n = 1 , nlnd_loc(localPet)+ nocn_loc(localPet)
530+ write (msg,' (A,2I8)' trim (subname)// ' : n, gindex_new = '
531+ call ESMF_LogWrite(trim (msg), ESMF_LOGMSG_INFO)
532+ end do
533+ end if
534+
535+ ! ---------------------
536+ ! Update mesh with new decomposition
537+ ! ---------------------
538+
539+ ! create new distgrid with new index
540+ distgrid_new = ESMF_DistGridCreate(arbSeqIndexList= gindex_new, rc= rc)
541+ if (chkerr(rc,__LINE__,u_FILE_u)) return
542+
543+ ! create new mesh with new distgrid
544+ mesh = ESMF_MeshCreate(noahmp% domain% mesh, elementDistGrid= distgrid_new, rc= rc)
545+ if (chkerr(rc,__LINE__,u_FILE_u)) return
546+
547+ ! destroy old and replace with new one
548+ call ESMF_MeshDestroy(noahmp% domain% mesh, rc= rc)
549+ if (chkerr(rc,__LINE__,u_FILE_u)) return
550+ noahmp% domain% mesh = mesh
551+
552+ ! ---------------------
553+ ! fix mask and fraction to be consistent with new decomposition
554+ ! ---------------------
555+
556+ ! mask
557+ noahmp% domain% mask(:nlnd_loc(localPet)) = 1
558+ noahmp% domain% mask(nlnd_loc(localPet)+ 1 :) = 0
559+
560+ ! fraction, read from file again
561+ allocate (tmpr4 (noahmp% domain% begl:noahmp% domain% endl))
562+ tmpr4 (:) = 0.0
563+ filename = trim (noahmp% nmlist% input_dir)// ' oro_data.tile*.nc'
564+ flds(1 )% short_name = ' land_frac'
565+ flds(1 )% ptr1r4 = > tmpr4
566+ call read_tiled_file(noahmp, filename, flds, rc= rc)
567+ if (ChkErr(rc,__LINE__,u_FILE_u)) return
568+ deallocate (tmpr4 )
569+
570+ ! ---------------------
571+ ! Clean memory
572+ ! ---------------------
573+
574+ deallocate (nlnd_loc)
575+ deallocate (nocn_loc)
576+ deallocate (mask_glb)
577+ deallocate (gindex_loc)
578+ deallocate (gindex_glb)
579+ deallocate (gindex_lnd)
580+ deallocate (gindex_ocn)
581+ deallocate (gindex_new)
582+ deallocate (lsize_arr)
583+
584+ call ESMF_LogWrite(subname// ' done'
585+
586+ end subroutine lnd_modify_decomp
587+
342588end module lnd_comp_domain
0 commit comments