Merge pull request #3055 from GEOS-ESM/feature/pchakrab/vertical-regr…

…idding

Vertical regridding - fixed-level to fixed-level

generic3g/actions/VerticalRegridActionNew.F90

@@ -0,0 +1,75 @@
+#include "MAPL_Generic.h"
+
+module mapl3g_VerticalRegridActionNew
+
+   use mapl_ErrorHandling
+   use mapl3g_ExtensionAction
+   use mapl3g_VerticalRegridMethod, only: VerticalRegridMethod_Flag
+   use mapl3g_CSR_SparseMatrix
+   use esmf
+   use, intrinsic :: iso_fortran_env, only: REAL32
+
+   implicit none
+   private
+
+   public :: VerticalRegridAction
+
+   type, extends(ExtensionAction) :: VerticalRegridAction
+      real(REAL32), allocatable :: src_vertical_coord(:)
+      real(REAL32), allocatable :: dst_vertical_coord(:)
+      type(VerticalRegridMethod_Flag) :: regrid_method
+      type(CSR_SparseMatrix_sp), allocatable :: weights(:) ! size of horz dims
+   contains
+      procedure :: initialize
+      procedure :: run
+      procedure, private :: compute_weights_
+   end type VerticalRegridAction
+
+   interface VerticalRegridAction
+      procedure :: new_VerticalRegridAction
+   end interface VerticalRegridAction
+
+contains
+
+   function new_VerticalRegridAction(src_vertical_coord, dst_vertical_coord, regrid_method) result(action)
+      type(VerticalRegridAction) :: action
+      real(REAL32), intent(in) :: src_vertical_coord(:)
+      real(REAL32), intent(in) :: dst_vertical_coord(:)
+      type(VerticalRegridMethod_Flag), intent(in) :: regrid_method
+
+      action%src_vertical_coord = src_vertical_coord
+      action%dst_vertical_coord = dst_vertical_coord
+
+      action%regrid_method = regrid_method
+   end function new_VerticalRegridAction
+
+   subroutine initialize(this, importState, exportState, clock, rc)
+      class(VerticalRegridAction), intent(inout) :: this
+      type(ESMF_State) :: importState
+      type(ESMF_State) :: exportState
+      type(ESMF_Clock) :: clock      
+      integer, optional, intent(out) :: rc
+
+      call this%compute_weights_()
+
+      _RETURN(_SUCCESS)
+   end subroutine initialize
+
+   subroutine run(this, importState, exportState, clock, rc)
+      class(VerticalRegridAction), intent(inout) :: this
+      type(ESMF_State) :: importState
+      type(ESMF_State) :: exportState
+      type(ESMF_Clock) :: clock      
+      integer, optional, intent(out) :: rc
+
+      ! call use_weights_to_compute_f_out_from_f_in()
+
+      _RETURN(_SUCCESS)
+   end subroutine run
+
+   subroutine compute_weights_(this)
+      class(VerticalRegridAction), intent(inout) :: this
+      ! this%weights = ...
+   end subroutine compute_weights_
+
+end module mapl3g_VerticalRegridActionNew

generic3g/tests/CMakeLists.txt

@@ -32,6 +32,7 @@ set (test_srcs
 
   Test_ModelVerticalGrid.pf
   Test_FixedLevelsVerticalGrid.pf
+  Test_VerticalLinearMap.pf
 
   Test_CSR_SparseMatrix.pf
   )

generic3g/tests/Test_VerticalLinearMap.pf

@@ -0,0 +1,45 @@
+#include "MAPL_TestErr.h"
+
+module Test_VerticalLinearMap
+
+   use mapl3g_CSR_SparseMatrix, only: SparseMatrix_sp => CSR_SparseMatrix_sp, matmul
+   use mapl3g_VerticalLinearMap, only: compute_linear_map_fixedlevels_to_fixedlevels
+   use funit
+   use, intrinsic :: iso_fortran_env, only: REAL32
+
+   implicit none
+
+contains
+
+   @test
+   subroutine test_linear_map_fixedlevels_to_fixedlevels()
+
+      real(REAL32), allocatable :: vcoord_src(:), vcoord_dst(:)
+      real(REAL32), allocatable :: fin(:)
+      ! real(REAL32), allocatable :: matrix(:, :)
+      type(SparseMatrix_sp) :: matrix
+      integer :: status
+
+      vcoord_src = [30., 20., 10.]
+      vcoord_dst = [20., 10.]
+      call compute_linear_map_fixedlevels_to_fixedlevels(vcoord_src, vcoord_dst, matrix, _RC)
+      fin = [7., 8., 3.]
+      @assertEqual([8., 3.], matmul(matrix, fin))
+
+      vcoord_src = [30., 20., 10.]
+      vcoord_dst = [25., 15.]
+      call compute_linear_map_fixedlevels_to_fixedlevels(vcoord_src, vcoord_dst, matrix, _RC)
+      fin = [2., 4., 6.]
+      @assertEqual([3.,5.], matmul(matrix, fin))
+      fin = [7., 8., 3.]
+      @assertEqual([7.5, 5.5], matmul(matrix, fin))
+
+      vcoord_src = [30., 20., 10.]
+      vcoord_dst = [28., 12.]
+      call compute_linear_map_fixedlevels_to_fixedlevels(vcoord_src, vcoord_dst, matrix, _RC)
+      fin = [20., 10., 5.]
+      @assertEqual([18., 6.], matmul(matrix, fin))
+
+   end subroutine test_linear_map_fixedlevels_to_fixedlevels
+
+end module Test_VerticalLinearMap

generic3g/vertical/CMakeLists.txt

@@ -5,7 +5,8 @@ target_sources(MAPL.generic3g PRIVATE
   MirrorVerticalGrid.F90
   FixedLevelsVerticalGrid.F90
   ModelVerticalGrid.F90
-
+  VerticalRegridMethod.F90
+  VerticalLinearMap.F90
   CSR_SparseMatrix.F90
 )
 
@@ -21,3 +22,8 @@ esma_add_fortran_submodules(
   SOURCES can_connect_to.F90
 )
 
+ecbuild_add_executable(
+  TARGET Test_VerticalLinearMap.x
+  SOURCES Test_VerticalLinearMap.F90
+  DEPENDS MAPL.generic3g ESMF::ESMF)
+target_link_libraries(Test_VerticalLinearMap.x PRIVATE ${this})

generic3g/vertical/Test_VerticalLinearMap.F90

@@ -0,0 +1,37 @@
+#define I_AM_MAIN
+#include "MAPL_Generic.h"
+
+program Test_VerticalLinearMap
+
+   use mapl_ErrorHandling
+   use mapl3g_CSR_SparseMatrix, only: SparseMatrix_sp => CSR_SparseMatrix_sp, matmul
+   use mapl3g_VerticalLinearMap, only: compute_linear_map_fixedlevels_to_fixedlevels
+   ! use mapl3g_VerticalLinearMap, only: apply_linear_map
+   use, intrinsic :: iso_fortran_env, only: REAL32
+
+   implicit none
+
+   real(REAL32), allocatable :: src(:), dst(:), fin(:)
+   ! real(REAL32), allocatable :: matrix(:, :)
+   type(SparseMatrix_sp) :: matrix
+   integer :: status
+
+   src = [30., 20., 10.]
+   dst = [20., 10.]
+   call compute_linear_map_fixedlevels_to_fixedlevels(src, dst, matrix, _RC)
+   fin = [7., 8., 3.]
+   print *, "Expected: [8.0, 3.0]", ", found: ", matmul(matrix, fin)
+
+   src = [30., 20., 10.]
+   dst = [25., 15.]
+   call compute_linear_map_fixedlevels_to_fixedlevels(src, dst, matrix, _RC)
+   fin = [7., 8., 3.]
+   print *, "Expected: [7.5, 5.5]", ", found: ", matmul(matrix, fin)
+
+   src = [30., 20., 10.]
+   dst = [28., 11.]
+   call compute_linear_map_fixedlevels_to_fixedlevels(src, dst, matrix, _RC)
+   fin = [7., 8., 3.]
+   print *, "Expected: [7.2, 3.5]", ", found: ", matmul(matrix, fin)
+
+end program Test_VerticalLinearMap

generic3g/vertical/VerticalLinearMap.F90

@@ -0,0 +1,118 @@
+#include "MAPL_Generic.h"
+
+module mapl3g_VerticalLinearMap
+
+   use mapl_ErrorHandling
+   use mapl3g_CSR_SparseMatrix, only: SparseMatrix_sp => CSR_SparseMatrix_sp
+   use mapl3g_CSR_SparseMatrix, only: add_row
+   use mapl3g_CSR_SparseMatrix, only: sparse_matmul_sp => matmul
+   use, intrinsic :: iso_fortran_env, only: REAL32
+
+   implicit none
+   private
+
+   public :: compute_linear_map_fixedlevels_to_fixedlevels
+
+   type IndexValuePair
+      integer :: index
+      real(REAL32) :: value_
+   end type IndexValuePair
+
+   interface operator(==)
+      procedure equal_to
+   end interface operator(==)
+
+   interface operator(/=)
+      procedure not_equal_to
+   end interface operator(/=)
+
+contains
+
+   ! Compute linear interpolation transformation matrix (src*matrix = dst)
+   ! when regridding (vertical) from fixed-levels to fixed-levels
+   ! NOTE: find_bracket_ below ASSUMEs that src array is monotonic and decreasing
+   subroutine compute_linear_map_fixedlevels_to_fixedlevels(src, dst, matrix, rc)
+      real(REAL32), intent(in) :: src(:)
+      real(REAL32), intent(in) :: dst(:)
+      type(SparseMatrix_sp), intent(out) :: matrix
+      ! real(REAL32), allocatable, intent(out) :: matrix(:, :)
+      integer, optional, intent(out) :: rc
+
+      real(REAL32) :: val, weight(2)
+      integer :: ndx, status
+      type(IndexValuePair) :: pair(2) ! [pair(1), pair(2)] is a bracket
+
+      _ASSERT(maxval(dst) <= maxval(src), "maxval(dst) > maxval(src)")
+      _ASSERT(minval(dst) >= minval(src), "minval(dst) < minval(src)")
+
+      ! allocate(matrix(size(dst), size(src)), source=0., _STAT)
+      ! Expected 2 non zero entries in each row
+      matrix = SparseMatrix_sp(size(dst), size(src), 2*size(dst))
+      do ndx = 1, size(dst)
+         val = dst(ndx)
+         call find_bracket_(val, src, pair)
+         call compute_linear_interpolation_weights_(val, pair%value_, weight)
+         if (pair(1) == pair(2)) then
+            ! matrix(ndx, pair(1)%index) = weight(1)
+            call add_row(matrix, ndx, pair(1)%index, [weight(1)])
+         else
+            ! matrix(ndx, pair(1)%index) = weight(1)
+            ! matrix(ndx, pair(2)%index) = weight(2)
+            call add_row(matrix, ndx, pair(1)%index, [weight(1), weight(2)])
+         end if
+      end do
+
+      _RETURN(_SUCCESS)
+   end subroutine compute_linear_map_fixedlevels_to_fixedlevels
+
+   ! Find array bracket containing val
+   ! ASSUME: array is monotonic and decreasing
+   subroutine find_bracket_(val, array, pair)
+      real(REAL32), intent(in) :: val
+      real(REAL32), intent(in) :: array(:)
+      Type(IndexValuePair), intent(out) :: pair(2)
+
+      integer :: ndx1, ndx2
+
+      ndx1 = minloc(abs(array - val), 1)
+      if (array(ndx1) < val) then
+         ndx1 = ndx1 - 1
+      end if
+      ndx2 = ndx1 ! array(ndx1) == val
+      if (array(ndx1) /= val) then
+         ndx2 = ndx1 +1
+      end if
+
+      pair(1) = IndexValuePair(ndx1, array(ndx1))
+      pair(2) = IndexValuePair(ndx2, array(ndx2))
+   end subroutine find_bracket_
+
+   subroutine compute_linear_interpolation_weights_(val, value_, weight)
+      real(REAL32), intent(in) :: val
+      real(REAL32), intent(in) :: value_(2)
+      real(REAL32), intent(out) :: weight(2)
+
+      real(REAL32) :: denominator, epsilon_sp
+
+      denominator = abs(value_(2) - value_(1))
+      epsilon_sp = epsilon(1.0)
+      if (denominator < epsilon_sp) then
+         weight = 1.0
+      else
+         weight(1) = abs(value_(2) - val)/denominator
+         weight(2) = abs(val - value_(1))/denominator
+      end if
+   end subroutine compute_linear_interpolation_weights_
+
+   elemental logical function equal_to(a, b)
+      type(IndexValuePair), intent(in) :: a, b
+      equal_to = .false.
+      equal_to = ((a%index == b%index) .and. (a%value_ == b%value_))
+   end function equal_to
+
+   elemental logical function not_equal_to(a, b)
+      type(IndexValuePair), intent(in) :: a, b
+      not_equal_to = .not. (a==b)
+   end function not_equal_to
+
+end module mapl3g_VerticalLinearMap

generic3g/vertical/VerticalRegridMethod.F90

@@ -0,0 +1,43 @@
+#include "MAPL_Generic.h"
+
+module mapl3g_VerticalRegridMethod
+
+   implicit none
+   private
+
+   public :: VerticalRegridMethod_Flag
+   public :: VERTICAL_REGRID_UNKNOWN
+   public :: VERTICAL_REGRID_LINEAR
+   public :: VERTICAL_REGRID_CONSERVATIVE
+   public :: operator(==), operator(/=)
+
+   type :: VerticalRegridMethod_Flag
+      private
+      integer :: id = -1
+   end type VerticalRegridMethod_Flag
+
+   interface operator(==)
+      procedure :: equal_to
+   end interface operator(==)
+
+   interface operator(/=)
+      procedure :: not_equal_to
+   end interface operator(/=)
+
+   type(VerticalRegridMethod_Flag), parameter :: VERTICAL_REGRID_UNKNOWN = VerticalRegridMethod_Flag(-1)
+   type(VerticalRegridMethod_Flag), parameter :: VERTICAL_REGRID_LINEAR = VerticalRegridMethod_Flag(1)
+   type(VerticalRegridMethod_Flag), parameter :: VERTICAL_REGRID_CONSERVATIVE = VerticalRegridMethod_Flag(2)
+
+contains
+
+   elemental logical function equal_to(a, b)
+      type(VerticalRegridMethod_Flag), intent(in) :: a, b
+      equal_to = (a%id == b%id)
+   end function equal_to
+
+   elemental logical function not_equal_to(a, b)
+      type(VerticalRegridMethod_Flag), intent(in) :: a, b
+      not_equal_to = .not. (a==b)
+   end function not_equal_to
+
+end module mapl3g_VerticalRegridMethod