Refactor surface restart logic in FV3GFS_io.F90 to not use hard-coded indices.

[SamuelTrahanNOAA]

Description

The surface restart uses hard-coded indices, which makes it difficult to add new restart variables, especially if different land surface schemes need different numbers of variables. This is problematic for developer @tanyasmirnova who has that exact problem.

The code now uses a variable nt (number of things) it increments before copying to or from a slice of a surface field array. Most of the copies happen in copy_to_GFS_Data or copy_from_GFS_Data, which wraps this into a tight, vectorizable, loop, while also making the syntax concise.

That means you'll see two constructs a lot.

Here, you see two columns: one increments nt, and one does the assignment. By having the two statements on one line, it is easy to spot if there's a missing nt=nt+1:

      nt=nt+1 ; sfc_name2(nt) = 'alvsf'
      nt=nt+1 ; sfc_name2(nt) = 'alvwf'
      nt=nt+1 ; sfc_name2(nt) = 'alnsf'
      nt=nt+1 ; sfc_name2(nt) = 'alnwf'
      nt=nt+1 ; sfc_name2(nt) = 'facsf'
      nt=nt+1 ; sfc_name2(nt) = 'facwf'
      nt=nt+1 ; sfc_name2(nt) = 'vfrac'
      nt=nt+1 ; sfc_name2(nt) = 'canopy'
      nt=nt+1 ; sfc_name2(nt) = 'f10m'
      nt=nt+1 ; sfc_name2(nt) = 't2m'
      nt=nt+1 ; sfc_name2(nt) = 'q2m'
      nt=nt+1 ; sfc_name2(nt) = 'vtype'
      nt=nt+1 ; sfc_name2(nt) = 'stype'

Most of the slices are copied like so:

        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%alvwf)   !--- alvwf                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%alnsf)   !--- alnsf                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%alnwf)   !--- alnwf                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%facsf)   !--- facsf                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%facwf)   !--- facwf                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%vfrac)   !--- vfrac                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%canopy)  !--- canopy                                                                                                                                                   
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%f10m)    !--- f10m                                                                                                                                                     
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%t2m)     !--- t2m                                                                                                                                                      
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%q2m)     !--- q2m                                                                                                                                                      
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%vtype)   !--- vtype                                                                                                                                                    
        call copy_to_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%stype)   !--- stype                                                                                                                                                    

where the nt=nt+1 and data copying logic is hidden behind the subroutine.

There were a few odd copies, with bizarre indexing (-2:4) or ones with calculations. For those, the logic is made explicitly clear:

            nt=nt+1
            do lsoil = -2, 4
              do ix = 1, Atm_block%blksz(nb)
                Sfcprop(nb)%zsnsoxy(ix,lsoil)  = sfc_var3zn(ii1(ix),jj1(ix),lsoil,nt)
              enddo
            enddo
...
       do k = 1,Model%kice
         do ix = 1, Atm_block%blksz(nb)
           ice=Sfcprop(nb)%tiice(ix,k)
           if(ice<one) then
             sfc_var3ice(ii1(ix),jj1(ix),k) = zero
           else
             sfc_var3ice(ii1(ix),jj1(ix),k) = ice
           endif
         enddo
       enddo

The code is also refactored to have the loops in an order that should be faster and easier to vectorize: nb, variable, k, ix.

!$omp parallel do default(shared) private(i, j, nb, ix, nt, ii1, jj1, lsoil, nt_save, k, ice)                                                                                                                                                
    block_loop: do nb = 1, Atm_block%nblks
       allocate(ii1(Atm_block%blksz(nb)))
       allocate(jj1(Atm_block%blksz(nb)))
       ii1=Atm_block%index(nb)%ii - isc + 1
       jj1=Atm_block%index(nb)%jj - jsc + 1
...
       call copy_from_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%albdifnir_lnd)
       call copy_from_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%emis_lnd)
       call copy_from_GFS_Data(ii1,jj1,isc,jsc,nt,sfc_var2,Sfcprop(nb)%emis_ice)
...
       do k = 1,Model%kice
         do ix = 1, Atm_block%blksz(nb)

When the code calls copy_to_GFS_Data or copy_from_GFS_Data, the k and ix loops are hidden behind subroutines:

   pure subroutine copy_from_GFS_Data_2d_phys2phys(ii1,jj1,isc,jsc,nt,var2d,var_block)
...
     nt=nt+1
     do ix=1,size(var_block)
       var2d(ii1(ix),jj1(ix),nt) = var_block(ix)
     enddo
   end subroutine copy_from_GFS_Data_2d_phys2phys

   pure subroutine copy_from_GFS_Data_3d_phys2phys(ii1,jj1,isc,jsc,nt,var3d,var_block)
...
     nt=nt+1
     do k=lbound(var_block,2),ubound(var_block,2)
       do ix=1,size(var_block,1)
         var3d(ii1(ix),jj1(ix),k,nt) = var_block(ix,k)
       enddo
     enddo
   end subroutine copy_from_GFS_Data_3d_phys2phys

Issue(s) addressed

fixes #604

Testing

hera.intel and hera.gnu ufs weather model regression tests

[joeolson42]

I'm not too familiar with this code, but the PR is well documented and the code looks clean.

[tanyasmirnova]

@SamuelTrahanNOAA Very useful modifications to FV3GFS_io.F90 to remove hard-wired numbers! This subroutine used to be a culprit when some variables have to be added to the input or restart file. Thank you very much, Sam, for your tremendous work to make this subroutine more flexible.

[DusanJovic-NOAA]

@SamuelTrahanNOAA Please update 'atmos_cubed_sphere' submodule to 16e659e5 hash from my branch.

[jkbk2004]

All tests are done on ufs-community/ufs-weather-model#1497. We can start merging in this pr.