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

…idding

Vertical regridding - step 3 - Model -> FixedLevels

generic3g/ComponentSpecParser.F90

@@ -1,6 +1,7 @@
 #include "MAPL_ErrLog.h"
 
 module mapl3g_ComponentSpecParser
+
    use mapl3g_ComponentSpec
    use mapl3g_ChildSpec
    use mapl3g_ChildSpecMap
@@ -23,12 +24,13 @@ module mapl3g_ComponentSpecParser
    use mapl3g_Stateitem
    use mapl3g_ESMF_Utilities
    use mapl3g_UserSetServices
+   use mapl3g_StateRegistry
    use gftl2_StringVector, only: StringVector
    use esmf
+
    implicit none
    private
 
-   !
    public :: parse_component_spec
 
    ! The following interfaces are public only for testing purposes.
@@ -63,15 +65,17 @@ module mapl3g_ComponentSpecParser
    !>
    ! Submodule declarations
    INTERFACE
-      module function parse_component_spec(hconfig, rc) result(spec)
+      module function parse_component_spec(hconfig, registry, rc) result(spec)
          type(ComponentSpec) :: spec
          type(ESMF_HConfig), target, intent(inout) :: hconfig
+         type(StateRegistry), optional, intent(in) :: registry
          integer, optional, intent(out) :: rc
       end function parse_component_spec
 
-      module function parse_geometry_spec(mapl_cfg, rc) result(geometry_spec)
+      module function parse_geometry_spec(mapl_cfg, registry, rc) result(geometry_spec)
          type(GeometrySpec) :: geometry_spec
          type(ESMF_HConfig), intent(in) :: mapl_cfg
+         type(StateRegistry), optional, intent(in) :: registry
          integer, optional, intent(out) :: rc
       end function parse_geometry_spec
 

generic3g/ComponentSpecParser/parse_component_spec.F90

@@ -4,9 +4,10 @@
 
 contains
 
-   module function parse_component_spec(hconfig, rc) result(spec)
+   module function parse_component_spec(hconfig, registry, rc) result(spec)
       type(ComponentSpec) :: spec
       type(ESMF_HConfig), target, intent(inout) :: hconfig
+      type(StateRegistry), optional, intent(in) :: registry
       integer, optional, intent(out) :: rc
 
       integer :: status
@@ -17,7 +18,7 @@ module function parse_component_spec(hconfig, rc) result(spec)
       _RETURN_UNLESS(has_mapl_section)
       mapl_cfg = ESMF_HConfigCreateAt(hconfig, keyString=MAPL_SECTION, _RC)
 
-      spec%geometry_spec = parse_geometry_spec(mapl_cfg, _RC)
+      spec%geometry_spec = parse_geometry_spec(mapl_cfg, registry, _RC)
       spec%var_specs = parse_var_specs(mapl_cfg, _RC)
       spec%connections = parse_connections(mapl_cfg, _RC)
       spec%children = parse_children(mapl_cfg, _RC)

generic3g/ComponentSpecParser/parse_geometry_spec.F90

@@ -1,18 +1,22 @@
 #include "MAPL_ErrLog.h"
 
 submodule (mapl3g_ComponentSpecParser) parse_geometry_spec_smod
+
    use mapl3g_VerticalGrid
    use mapl3g_BasicVerticalGrid
    use mapl3g_FixedLevelsVerticalGrid
+   use mapl3g_ModelVerticalGrid
+
    implicit none(external,type)
 
 contains
 
    ! Geom subcfg is passed raw to the GeomManager layer.  So little
    ! processing is needed here.
-   module function parse_geometry_spec(mapl_cfg, rc) result(geometry_spec)
+   module function parse_geometry_spec(mapl_cfg, registry, rc) result(geometry_spec)
       type(GeometrySpec) :: geometry_spec
       type(ESMF_HConfig), intent(in) :: mapl_cfg
+      type(StateRegistry), optional, intent(in) :: registry
       integer, optional, intent(out) :: rc
 
       integer :: status
@@ -29,7 +33,7 @@ module function parse_geometry_spec(mapl_cfg, rc) result(geometry_spec)
       type(GeomManager), pointer :: geom_mgr
       class(GeomSpec), allocatable :: geom_spec
       integer :: num_levels
-      character(:), allocatable :: vertical_grid_class, standard_name, units
+      character(:), allocatable :: vertical_grid_class, standard_name, units, short_name
       class(VerticalGrid), allocatable :: vertical_grid
       real, allocatable :: levels(:)
 
@@ -102,6 +106,15 @@ module function parse_geometry_spec(mapl_cfg, rc) result(geometry_spec)
             units = ESMF_HConfigAsString(vertical_grid_cfg, keyString='units', _RC)
             levels = ESMF_HConfigAsR4Seq(vertical_grid_cfg, keyString='levels' ,_RC)
             vertical_grid = FixedLevelsVerticalGrid(standard_name, levels, units)
+         case('model')
+            num_levels = ESMF_HConfigAsI4(vertical_grid_cfg, keyString='num_levels', _RC)
+            vertical_grid = ModelVerticalGrid(num_levels=num_levels)
+            short_name = ESMF_HConfigAsString(vertical_grid_cfg, keyString='short_name', _RC)
+            select type(vertical_grid)
+            type is(ModelVerticalGrid)
+               call vertical_grid%add_variant(short_name=short_name)
+               call vertical_grid%set_registry(registry)
+            end select
          case default
             _FAIL('vertical grid class '//vertical_grid_class//' not supported')
          end select
@@ -112,4 +125,3 @@ module function parse_geometry_spec(mapl_cfg, rc) result(geometry_spec)
    end function parse_geometry_spec
 
 end submodule parse_geometry_spec_smod
-

generic3g/OuterMetaComponent/SetServices.F90

@@ -32,7 +32,7 @@ recursive module subroutine SetServices_(this, rc)
       integer :: status
       type(ESMF_GridComp) :: user_gridcomp
 
-      this%component_spec = parse_component_spec(this%hconfig, _RC)
+      this%component_spec = parse_component_spec(this%hconfig, this%registry, _RC)
       user_gridcomp = this%user_gc_driver%get_gridcomp()
       call attach_inner_meta(user_gridcomp, this%self_gridcomp, _RC)
       call this%user_setservices%run(user_gridcomp, _RC)

generic3g/actions/VerticalRegridAction.F90

@@ -8,7 +8,7 @@ module mapl3g_VerticalRegridAction
    use mapl3g_CouplerPhases, only: GENERIC_COUPLER_UPDATE
    use mapl3g_VerticalRegridMethod
    use mapl3g_VerticalLinearMap, only: compute_linear_map
-   use mapl3g_CSR_SparseMatrix, only: SparseMatrix_sp => CSR_SparseMatrix_sp, matmul
+   use mapl3g_CSR_SparseMatrix, only: SparseMatrix_sp => CSR_SparseMatrix_sp, matmul, shape
    use mapl3g_FieldCondensedArray, only: assign_fptr_condensed_array
    use esmf
 
@@ -23,7 +23,7 @@ module mapl3g_VerticalRegridAction
 
    type, extends(ExtensionAction) :: VerticalRegridAction
       type(ESMF_Field) :: v_in_coord, v_out_coord
-      type(SparseMatrix_sp) :: matrix
+      type(SparseMatrix_sp), allocatable :: matrix(:, :)
       type(GriddedComponentDriver), pointer :: v_in_coupler => null()
       type(GriddedComponentDriver), pointer :: v_out_coupler => null()
       type(VerticalRegridMethod) :: method = VERTICAL_REGRID_UNKNOWN
@@ -65,9 +65,9 @@ subroutine initialize(this, importState, exportState, clock, rc)
       type(ESMF_Clock) :: clock
       integer, optional, intent(out) :: rc
 
-      real(ESMF_KIND_R4), pointer :: vcoord_in(:)
-      real(ESMF_KIND_R4), pointer :: vcoord_out(:)
-      integer :: status
+      real(ESMF_KIND_R4), pointer :: v_in(:, :, :), v_out(:, :, :)
+      integer :: shape_in(3), shape_out(3), n_horz, n_ungridded
+      integer :: horz1, horz2, ungrd, status
 
       _ASSERT(this%method == VERTICAL_REGRID_LINEAR, "regrid method can only be linear")
 
@@ -79,10 +79,28 @@ subroutine initialize(this, importState, exportState, clock, rc)
       !    call this%v_out_coupler%initialize(_RC)
       ! end if
 
-      call ESMF_FieldGet(this%v_in_coord, fArrayPtr=vcoord_in, _RC)
-      call ESMF_FieldGet(this%v_out_coord, fArrayPtr=vcoord_out, _RC)
-
-      call compute_linear_map(vcoord_in, vcoord_out, this%matrix, RC)
+      call assign_fptr_condensed_array(this%v_in_coord, v_in, _RC)
+      shape_in = shape(v_in)
+      n_horz = shape_in(1)
+      n_ungridded = shape_in(3)
+
+      call assign_fptr_condensed_array(this%v_out_coord, v_out, _RC)
+      shape_out = shape(v_out)
+      _ASSERT((shape_in(1) == shape_out(1)), "horz dims are expected to be equal")
+      _ASSERT((shape_in(3) == shape_out(3)), "ungridded dims are expected to be equal")
+
+      allocate(this%matrix(n_horz, n_horz))
+
+      ! TODO: Convert to a `do concurrent` loop
+      do horz1 = 1, n_horz
+         do horz2 = 1, n_horz
+            do ungrd = 1, n_ungridded
+               associate(src => v_in(horz1, :, ungrd), dst => v_out(horz2, :, ungrd))
+                 call compute_linear_map(src, dst, this%matrix(horz1, horz2), _RC)
+               end associate
+            end do
+         end do
+      end do
 
       _RETURN(_SUCCESS)
    end subroutine initialize
@@ -98,7 +116,8 @@ subroutine update(this, importState, exportState, clock, rc)
       integer :: status
       type(ESMF_Field) :: f_in, f_out
       real(ESMF_KIND_R4), pointer :: x_in(:,:,:), x_out(:,:,:)
-      integer :: x_shape(3), horz, ungridded
+      integer :: shape_in(3), shape_out(3), n_horz, n_ungridded
+      integer :: horz1, horz2, ungrd
 
       ! if (associated(this%v_in_coupler)) then
       !    call this%v_in_coupler%run(phase_idx=GENERIC_COUPLER_UPDATE, _RC)
@@ -110,13 +129,19 @@ subroutine update(this, importState, exportState, clock, rc)
 
       call ESMF_StateGet(importState, itemName='import[1]', field=f_in, _RC)
       call assign_fptr_condensed_array(f_in, x_in, _RC)
+      shape_in = shape(x_in)
+      n_horz = shape_in(1)
+      n_ungridded = shape_in(3)
 
       call ESMF_StateGet(exportState, itemName='export[1]', field=f_out, _RC)
       call assign_fptr_condensed_array(f_out, x_out, _RC)
+      shape_out = shape(x_out)
+
+      _ASSERT((shape_in(1) == shape_out(1)), "horz dims are expected to be equal")
+      _ASSERT((shape_in(3) == shape_out(3)), "ungridded dims are expected to be equal")
 
-      x_shape = shape(x_out)
-      do concurrent (horz=1:x_shape(1), ungridded=1:x_shape(3))
-         x_out(horz, :, ungridded) = matmul(this%matrix, x_in(horz, :, ungridded))
+      do concurrent (horz1=1:n_horz, horz2=1:n_horz, ungrd=1:n_ungridded)
+         x_out(horz2, :, ungrd) = matmul(this%matrix(horz1, horz2), x_in(horz1, :, ungrd))
       end do
 
       _RETURN(_SUCCESS)

generic3g/registry/StateItemExtension.F90

@@ -103,7 +103,7 @@ end subroutine add_consumer
    ! is added to the export specs of source (this), and the new extension
    ! gains it as a reference (pointer).
 
-   function make_extension(this, goal, rc) result(extension)
+   recursive function make_extension(this, goal, rc) result(extension)
       type(StateItemExtension), target :: extension
       class(StateItemExtension), target, intent(inout) :: this
       class(StateItemSpec), target, intent(in) :: goal

generic3g/registry/StateRegistry.F90

@@ -797,7 +797,7 @@ end function get_import_couplers
 
    ! Repeatedly extend family at v_pt until extension can directly
    ! connect to goal_spec.
-   function extend(registry, v_pt, goal_spec, rc) result(extension)
+   recursive function extend(registry, v_pt, goal_spec, rc) result(extension)
       use mapl3g_MultiState
       use mapl3g_ActualConnectionPt, only: ActualConnectionPt
       type(StateItemExtension), pointer :: extension

generic3g/specs/FieldSpec.F90

@@ -113,6 +113,9 @@ module mapl3g_FieldSpec
 
       procedure :: set_info
       procedure :: set_geometry
+
+      procedure :: write_formatted
+      generic :: write(formatted) => write_formatted
    end type FieldSpec
 
    interface FieldSpec
@@ -151,6 +154,7 @@ module mapl3g_FieldSpec
       type(ESMF_Geom), allocatable :: geom
       type(ESMF_TypeKind_Flag) :: typekind
       character(:), allocatable :: units
+      type(VerticalDimSpec), allocatable :: vertical_dim_spec
       type(VerticalRegridMethod), allocatable :: regrid_method
    contains
       procedure :: adapt_one => adapt_vertical_grid
@@ -329,6 +333,31 @@ subroutine allocate(this, rc)
       _RETURN(ESMF_SUCCESS)
    end subroutine allocate
 
+   subroutine write_formatted(this, unit, iotype, v_list, iostat, iomsg)
+      class(FieldSpec), intent(in) :: this
+      integer, intent(in) :: unit
+      character(*), intent(in) :: iotype
+      integer, intent(in) :: v_list(:)
+      integer, intent(out) :: iostat
+      character(*), intent(inout) :: iomsg
+
+      write(unit, "(a, a)", iostat=iostat, iomsg=iomsg) "FieldSpec(", new_line("a")
+      if (allocated(this%standard_name)) then
+         write(unit, "(3x, a, a, a)", iostat=iostat, iomsg=iomsg) "standard name:", this%standard_name, new_line("a")
+      end if
+      if (allocated(this%long_name)) then
+         write(unit, "(3x, a, a, a)", iostat=iostat, iomsg=iomsg) "long name:", this%long_name, new_line("a")
+      end if
+      if (allocated(this%units)) then
+         write(unit, "(3x, a, a, a)", iostat=iostat, iomsg=iomsg) "unit:", this%units, new_line("a")
+      end if
+      write(unit, "(3x, dt'g0', a)", iostat=iostat, iomsg=iomsg) this%vertical_dim_spec, new_line("a")
+      if (allocated(this%vertical_grid)) then
+         write(unit, "(3x, dt'g0', a)", iostat=iostat, iomsg=iomsg) this%vertical_grid, new_line("a")
+      end if
+      write(unit, "(a)") ")"
+   end subroutine write_formatted
+
    function get_ungridded_bounds(this, rc) result(bounds)
       type(LU_Bound), allocatable :: bounds(:)
       type(FieldSpec), intent(in) :: this
@@ -816,18 +845,20 @@ logical function adapter_match_geom(this, spec, rc) result(match)
       _RETURN(_SUCCESS)
    end function adapter_match_geom
 
-   function new_VerticalGridAdapter(vertical_grid, geom, typekind, units, regrid_method) result(vertical_grid_adapter)
+   function new_VerticalGridAdapter(vertical_grid, geom, typekind, units, vertical_dim_spec, regrid_method) result(vertical_grid_adapter)
       type(VerticalGridAdapter) :: vertical_grid_adapter
       class(VerticalGrid), optional, intent(in) :: vertical_grid
       type(ESMF_Geom), optional, intent(in) :: geom
       type(ESMF_Typekind_Flag), intent(in) :: typekind
       character(*), optional, intent(in) :: units
+      type(VerticalDimSpec), intent(in) :: vertical_dim_spec
       type(VerticalRegridMethod), optional, intent(in) :: regrid_method
 
       if (present(vertical_grid)) vertical_grid_adapter%vertical_grid = vertical_grid
       if (present(geom)) vertical_grid_adapter%geom = geom
       vertical_grid_adapter%typekind = typekind
       if (present(units)) vertical_grid_adapter%units = units
+      vertical_grid_adapter%vertical_dim_spec = vertical_dim_spec
       if (present(regrid_method)) vertical_grid_adapter%regrid_method = regrid_method
    end function new_VerticalGridAdapter
 
@@ -845,9 +876,9 @@ subroutine adapt_vertical_grid(this, spec, action, rc)
       select type (spec)
       type is (FieldSpec)
          call spec%vertical_grid%get_coordinate_field(v_in_coord, v_in_coupler, &
-              'ignore', spec%geom, spec%typekind, spec%units, _RC)
+              'ignore', spec%geom, spec%typekind, spec%units, spec%vertical_dim_spec, _RC)
          call this%vertical_grid%get_coordinate_field(v_out_coord, v_out_coupler, &
-              'ignore', this%geom, this%typekind, this%units, _RC)
+              'ignore', this%geom, this%typekind, this%units, this%vertical_dim_spec, _RC)
          action = VerticalRegridAction(v_in_coord, v_out_coupler, v_out_coord, v_out_coupler, this%regrid_method)
          spec%vertical_grid = this%vertical_grid
       end select
@@ -878,8 +909,9 @@ logical function same_vertical_grid(src_grid, dst_grid, rc)
          class(VerticalGrid), allocatable, intent(in) :: dst_grid
          integer, optional, intent(out) :: rc
 
-         same_vertical_grid = .true.
+         same_vertical_grid = .false.
          if (.not. allocated(dst_grid)) then
+            same_vertical_grid = .true.
             _RETURN(_SUCCESS) ! mirror grid
          end if
 
@@ -901,10 +933,11 @@ logical function same_vertical_grid(src_grid, dst_grid, rc)
             type is(FixedLevelsVerticalGrid)
                same_vertical_grid = (src_grid == dst_grid)
             class default
-               _FAIL("not implemented yet")
+               same_vertical_grid = .false.
             end select
          class default
-            _FAIL("not implemented yet")
+            same_vertical_grid = .false.
+            ! _FAIL("not implemented yet")
          end select
 
          _RETURN(_SUCCESS)
@@ -992,14 +1025,21 @@ recursive function make_adapters(this, goal_spec, rc) result(adapters)
       class(StateItemSpec), intent(in) :: goal_spec
       integer, optional, intent(out) :: rc
 
+      type(VerticalGridAdapter) :: vertical_grid_adapter
       integer :: status
 
       select type (goal_spec)
       type is (FieldSpec)
          allocate(adapters(4))
          allocate(adapters(1)%adapter, source=GeomAdapter(goal_spec%geom, goal_spec%regrid_param))
-         allocate(adapters(2)%adapter, &
-              source=VerticalGridAdapter(goal_spec%vertical_grid, goal_spec%geom, goal_spec%typekind, goal_spec%units, VERTICAL_REGRID_LINEAR))
+         vertical_grid_adapter = VerticalGridAdapter( &
+              goal_spec%vertical_grid, &
+              goal_spec%geom, &
+              goal_spec%typekind, &
+              goal_spec%units, &
+              goal_spec%vertical_dim_spec, &
+              VERTICAL_REGRID_LINEAR)
+         allocate(adapters(2)%adapter, source=vertical_grid_adapter)
          allocate(adapters(3)%adapter, source=TypeKindAdapter(goal_spec%typekind))
          allocate(adapters(4)%adapter, source=UnitsAdapter(goal_spec%units))
       type is (WildCardSpec)