Initial commit of hflux regridding improvement from GMAO

Signed-off-by: Lizzie Lundgren <[email protected]>

base/Base/Base_Base_implementation.F90

@@ -833,7 +833,7 @@ module subroutine MAPL_MakeDecomposition(nx, ny, unusable, reduceFactor, rc)
 
     type (ESMF_VM) :: vm
     integer :: pet_count
-
+    integer :: bias
     character(len=*), parameter :: Iam= __FILE__ // '::MAPL_MakeDecomposition()'
     integer :: status
 
@@ -843,11 +843,18 @@ module subroutine MAPL_MakeDecomposition(nx, ny, unusable, reduceFactor, rc)
     _VERIFY(status)
     call ESMF_VMGet(vm, petCount=pet_count, rc=status)
     _VERIFY(status)
-    if (present(reduceFactor)) pet_count=pet_count/reduceFactor
+    if (present(reduceFactor)) then
+       pet_count=pet_count/reduceFactor
+       ! Assume CS
+       bias = 1
+    else
+       ! Assume Lat-Lon
+       bias =2
+    end if
 
     ! count down from sqrt(n)
     ! Note: inal iteration (nx=1) is guaranteed to succeed.
-    do nx = floor(sqrt(real(2*pet_count))), 1, -1
+    do nx = nint(sqrt(real(bias*pet_count))), 1, -1
        if (mod(pet_count, nx) == 0) then ! found a decomposition
           ny = pet_count / nx
           exit

base/HorizontalFluxRegridder.F90

@@ -9,7 +9,9 @@ module mapl_HorizontalFluxRegridder
    use mapl_RegridMethods
    use mapl_KeywordEnforcerMod
    use mapl_ErrorHandlingMod
-   use mapl_BaseMod
+   use mapl_MaplGrid
+   use mapl_Base
+   use mapl_SphericalGeometry
    implicit none
    private
 
@@ -20,6 +22,8 @@ module mapl_HorizontalFluxRegridder
       integer :: resolution_ratio = -1
       integer :: im_in, jm_in
       integer :: im_out, jm_out
+      real, allocatable :: dx_in(:,:), dy_in(:,:)
+      real, allocatable :: dx_out(:,:), dy_out(:,:)
    contains
       procedure, nopass :: supports
       procedure :: initialize_subclass
@@ -54,6 +58,7 @@ logical function supports(spec, unusable, rc)
       call MAPL_GridGet(spec%grid_out, localCellCountPerDim=counts_out, __RC__)
 
       supports = all(mod(counts_in(1:2), counts_out(1:2)) == 0) .or. all(mod(counts_out, counts_in) == 0)
+      _ASSERT(supports, "HFlux regridder requires local domains to be properly nested.")
 
       _RETURN(_SUCCESS)
    end function supports
@@ -70,6 +75,8 @@ subroutine initialize_subclass(this, unusable, rc)
 
      integer :: counts(5)
      integer :: status
+     integer :: units ! unused
+     real(kind=ESMF_KIND_R8), allocatable :: corner_lons(:,:), corner_lats(:,:)
 
      _UNUSED_DUMMY(unusable)
      spec = this%get_spec()
@@ -91,6 +98,35 @@ subroutine initialize_subclass(this, unusable, rc)
          _ASSERT((IM_in / IM_out) == (JM_in / JM_out), 'inconsistent aspect ratio')
 
          this%resolution_ratio = (IM_in / IM_out)
+
+         allocate(corner_lons(IM_in+1,JM_in+1), corner_lats(IM_in+1,JM_in+1))
+         associate(lons => corner_lons, lats => corner_lats)
+           call MAPL_GridGetCorners(grid_in, gridCornerLons=lons, gridCornerLats=lats, _RC)
+
+           this%dx_in = distance( &
+                lons(1:IM_in,1:JM_in), lats(1:IM_in,1:JM_in), &
+                lons(2:IM_in+1,1:JM_in), lats(2:IM_in+1,1:JM_in))
+
+           this%dy_in = distance( &
+                lons(1:IM_in,1:JM_in), lats(1:IM_in,1:JM_in), &
+                lons(1:IM_in,2:JM_in+1), lats(1:IM_in,2:JM_in+1))
+         end associate
+
+         deallocate(corner_lons, corner_lats)
+         allocate(corner_lons(IM_out+1,JM_out+1), corner_lats(IM_out+1,JM_out+1))
+         associate(lons => corner_lons, lats => corner_lats)
+           call MAPL_GridGetCorners(grid_out, gridCornerLons=lons, gridCornerLats=lats, _RC)
+
+           this%dx_out = distance( &
+                lons(1:IM_in,1:JM_in), lats(1:IM_in,1:JM_in), &
+                lons(2:IM_in+1,1:JM_in), lats(2:IM_in+1,1:JM_in))
+
+           this%dy_out = distance( &
+                lons(1:IM_in,1:JM_in), lats(1:IM_in,1:JM_in), &
+                lons(1:IM_in,2:JM_in+1), lats(1:IM_in,2:JM_in+1))
+
+         end associate
+
        end associate
      end associate
 
@@ -129,9 +165,14 @@ subroutine regrid_vector_2d_real32(this, u_in, v_in, u_out, v_out, rotate, rc)
              do i  = 1, IM
                 m_y = 0
                 do ii = 1 + (i-1)*N, i*N
-                   m_y = m_y + v_in(ii,jj)
+                   associate (d_in => this%dx_in(ii,jj))
+                     m_y = m_y + v_in(ii,jj) * d_in
+                   end associate
                 end do
-                v_out(i,j) = m_y
+
+                associate (d_out => this%dx_out(i,j))
+                  v_out(i,j) = m_y / d_out
+                end associate
              end do
           end do
 
@@ -141,9 +182,13 @@ subroutine regrid_vector_2d_real32(this, u_in, v_in, u_out, v_out, rotate, rc)
              do j  = 1, JM
                 m_x = 0
                 do jj = 1 + (j-1)*N, j*N
-                   m_x = m_x + u_in(ii,jj)
+                   associate (d_in => this%dy_in(ii,jj))
+                     m_x = m_x + u_in(ii,jj) * d_in
+                   end associate
                 end do
-                u_out(i,j) = m_x
+                associate (d_out => this%dy_out(i,j))
+                  u_out(i,j) = m_x / d_out
+                end associate
              end do
           end do
 
@@ -186,9 +231,13 @@ subroutine regrid_vector_2d_real64(this, u_in, v_in, u_out, v_out, rotate, rc)
              do i  = 1, IM
                 m_y = 0
                 do ii = 1 + (i-1)*N, i*N
-                   m_y = m_y + v_in(ii,jj)
+                   associate (d_in => this%dx_in(ii,jj))
+                     m_y = m_y + v_in(ii,jj) * d_in
+                   end associate
                 end do
-                v_out(i,j) = m_y
+                associate (d_out => this%dx_out(i,j))
+                  v_out(i,j) = m_y / d_out
+                end associate
              end do
           end do
 
@@ -198,9 +247,13 @@ subroutine regrid_vector_2d_real64(this, u_in, v_in, u_out, v_out, rotate, rc)
              do j  = 1, JM
                 m_x = 0
                 do jj = 1 + (j-1)*N, j*N
-                   m_x = m_x + u_in(ii,jj)
+                   associate (d_in => this%dy_in(ii,jj))
+                     m_x = m_x + u_in(ii,jj) * d_in
+                   end associate
                 end do
-                u_out(i,j) = m_x
+                associate (d_out => this%dy_out(i,j))
+                  u_out(i,j) = m_x / d_out
+                end associate
              end do
           end do
 

base/MAPL_SphericalGeometry.F90

@@ -6,10 +6,19 @@ module MAPL_SphericalGeometry
    use MAPL_Constants
    use, intrinsic :: iso_fortran_env, only: REAL64,REAL32
 
-implicit none
-private
-public get_points_in_spherical_domain
-contains
+   implicit none
+   private
+   public :: get_points_in_spherical_domain
+   public :: distance
+
+
+   interface distance
+      module procedure distance_r32
+      module procedure distance_r64
+   end interface distance
+
+
+ contains
 
  subroutine get_points_in_spherical_domain(center_lons,center_lats,corner_lons,corner_lats,lons,lats,ii,jj,rc)
     real(real64), intent(in) :: center_lats(:,:),center_lons(:,:)
@@ -217,4 +226,29 @@ function vect_mag(v) result(mag)
    mag = sqrt(v(1)*v(1)+v(2)*v(2)+v(3)*v(3))
 end function vect_mag
 
+
+   ! Computes distance between two points (in lat lon as radians),
+   ! and returns distance in radians (unit sphere)
+   ! Using formulae from: https://www.movable-type.co.uk/scripts/latlong.html
+
+   elemental real(kind=REAL64) function distance_r64(lon1, lat1, lon2, lat2) result(d)
+      real(kind=REAL64), intent(in) :: lon1, lat1
+      real(kind=REAL64), intent(in) :: lon2, lat2
+
+      associate(a => sin(lat2-lat1)**2 + cos(lat1)*cos(lat2)*sin((lon2-lon1)/2)**2)
+        d = 2*atan2(sqrt(a), sqrt(1-a))
+      end associate
+
+   end function distance_r64
+
+   elemental real(kind=REAL32) function distance_r32(lon1, lat1, lon2, lat2) result(d)
+      real(kind=REAL32), intent(in) :: lon1, lat1
+      real(kind=REAL32), intent(in) :: lon2, lat2
+
+      associate(a => sin(lat2-lat1)**2 + cos(lat1)*cos(lat2)*sin((lon2-lon1)/2)**2)
+        d = 2*atan2(sqrt(a), sqrt(1-a))
+      end associate
+
+   end function distance_r32
+
 end module MAPL_SphericalGeometry