fv3jedi_vc_vertremap_mod.f90

Go to the documentation of this file.

! (C) Copyright 2020 UCAR
!
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 
module fv3jedi_vc_vertremap_mod
 
! fckit
use fckit_configuration_module, only: fckit_configuration
 
! fms
use constants_mod,         only: grav
use field_manager_mod,     only: model_atmos
use mpp_domains_mod,       only: mpp_update_domains
use tracer_manager_mod,    only: get_number_tracers, get_tracer_names, get_tracer_index, no_tracer, &
                                 set_tracer_profile
 
! fv3
use external_ic_mod,       only: remap_scalar, remap_dwinds, source
use fv_arrays_mod,         only: fv_atmos_type, deallocate_fv_atmos_type, r_grid
use fv_grid_utils_mod,     only: mid_pt_sphere, get_unit_vect2, get_latlon_vector, inner_prod
use test_cases_mod,        only: checker_tracers
 
! fv3jedi
use fv_prec_mod,           only: kind_fv3
use fv_init_mod,           only: fv_init
use fv3jedi_geom_mod,      only: fv3jedi_geom
use fv3jedi_fieldfail_mod, only: field_fail
use fv3jedi_field_mod,     only: copy_subset, field_clen, fv3jedi_field
use fv3jedi_kinds_mod,     only: kind_real
use fv3jedi_state_mod,     only: fv3jedi_state
 
implicit none
private
public :: fv3jedi_vc_vertremap
 
type :: fv3jedi_vc_vertremap
 type(fv_atmos_type), allocatable :: atm(:)
 logical, allocatable :: grids_on_this_pe(:)
 logical :: from_cold_start, checker_tr
 integer :: nt_checker
 contains
   procedure :: create
   procedure :: delete
   procedure :: changevar
end type fv3jedi_vc_vertremap
 
! --------------------------------------------------------------------------------------------------
 
contains
 
! --------------------------------------------------------------------------------------------------
 
subroutine create(self, geom, conf)
 
class(fv3jedi_vc_vertremap), intent(inout) :: self
type(fv3jedi_geom),          intent(in)    :: geom
type(fckit_configuration),   intent(in)    :: conf
 
integer :: gtile, p_split = 1, n
logical :: checks_passed
character(len=:), allocatable :: str
 
! Create Atm structure
call fv_init(self%Atm, 300.0_kind_real, self%grids_on_this_pe, p_split, gtile)
 
! Flag to use cold starts
if( .not. conf%get('input is cold starts', self%from_cold_start) ) self%from_cold_start = .true.
 
! Check tracers
if( .not. conf%get('check tracers',    self%checker_tr) ) self%checker_tr = .false.
if( .not. conf%get('check tracers nt', self%nt_checker) ) self%nt_checker = 0
 
! Set global source variable in external_ic
if (.not. conf%get("source of inputs", str)) then
  str = 'FV3GFS GAUSSIAN NETCDF FILE'
endif
source = trim(str)
 
! Remapping needs nggps_ic to be true
self%Atm(1)%flagstruct%nggps_ic = .true.
 
! ak/bk
self%Atm(1)%ak = real(geom%ak,kind_fv3)
self%Atm(1)%bk = real(geom%bk,kind_fv3)
self%Atm%ptop = real(geom%ak(1),kind_fv3)
 
! Sanity checks
checks_passed = .true.
if (checks_passed) checks_passed = geom%npx == self%Atm(1)%npx
if (checks_passed) checks_passed = geom%npy == self%Atm(1)%npy
if (checks_passed) checks_passed = geom%npz == self%Atm(1)%npz
if (checks_passed) checks_passed = geom%isd == self%Atm(1)%bd%isd
if (checks_passed) checks_passed = geom%ied == self%Atm(1)%bd%ied
if (checks_passed) checks_passed = geom%jsd == self%Atm(1)%bd%jsd
if (checks_passed) checks_passed = geom%jed == self%Atm(1)%bd%jed
if (.not.checks_passed) call abor1_ftn("fv3jedi_vc_vertremap_mod.field_fail: Geometry generated"// &
                                       " here does not match fv3-jedi geometry.")
 
if (.not.size(self%Atm)==1) call abor1_ftn("fv3jedi_vc_vertremap_mod.field_fail: Atm strucutre"// &
                                           " with size > 1 not supported.")
 
end subroutine create
 
! --------------------------------------------------------------------------------------------------
 
subroutine delete(self)
 
class(fv3jedi_vc_vertremap), intent(inout) :: self
 
integer :: n
 
do n = 1,size(self%Atm)
  call deallocate_fv_atmos_type(self%Atm(n))
enddo
deallocate(self%Atm)
deallocate(self%grids_on_this_pe)
 
end subroutine delete
 
! --------------------------------------------------------------------------------------------------
 
subroutine changevar(self, xin, xout)
 
class(fv3jedi_vc_vertremap), target, intent(inout) :: self
type(fv3jedi_state),                 intent(in)    :: xin
type(fv3jedi_state),                 intent(inout) :: xout
 
! Copy versus transform control
integer :: f
character(len=field_clen), allocatable :: fields_to_do(:)
type(fv3jedi_field), pointer :: field_ptr
 
! Local versions of fv3jedi state
real(kind=kind_real), allocatable :: orog_filt(:,:,:)
real(kind=kind_real), allocatable :: ps_cold(:,:,:)
real(kind=kind_real), allocatable :: zh_cold(:,:,:)
real(kind=kind_real), allocatable ::  w_cold(:,:,:)
real(kind=kind_real), allocatable ::  t_cold(:,:,:)
real(kind=kind_real), allocatable :: q_tmp(:,:,:)
real(kind=kind_real), allocatable :: ud_cold(:,:,:)
real(kind=kind_real), allocatable :: vd_cold(:,:,:)
 
logical :: have_remapped
type(fv_atmos_type), pointer :: Atm
integer:: i, j, k, nt, ntracers, ntprog, itoa, levp, isc, iec, jsc, jec, npz, nts
integer:: liq_wat, ice_wat, rainwat, snowwat, graupel, ntclamt
character(len=64) :: tracer_name
real(kind=kind_fv3), allocatable :: ak(:), bk(:)
real(kind=kind_fv3), allocatable ::  q_cold(:,:,:,:)
real(kind=kind_fv3) :: wt, qt, m_fac
 
! Field names
character(len=field_clen) :: ps_fname
character(len=field_clen) :: zh_fname
character(len=field_clen) :: w_fname
character(len=field_clen) :: t_fname
character(len=field_clen) :: ud_fname
character(len=field_clen) :: vd_fname
 
 
! Identity part of the change of variables
! ----------------------------------------
call copy_subset(xin%fields, xout%fields, fields_to_do)
 
 
! If variable change is the identity early exit
! ---------------------------------------------
if (.not.allocated(fields_to_do)) return
 
 
! Field names
! -----------
ps_fname = 'ps'
zh_fname = 'zh'
w_fname = 'w'
t_fname = 't'
ud_fname = 'ud'
vd_fname = 'vd'
if (self%from_cold_start) then
  ps_fname = trim(ps_fname)//'_cold'
  zh_fname = trim(zh_fname)//'_cold'
  w_fname = trim(w_fname)//'_cold'
  t_fname = trim(t_fname)//'_cold'
  ud_fname = trim(ud_fname)//'_cold'
  vd_fname = trim(vd_fname)//'_cold'
endif
 
 
! Remap to new Lagrangian vertical coordinate
! -------------------------------------------
have_remapped = .false.
if ( xin%has_field(ps_fname) .and. xin%has_field(zh_fname) .and. &
     xin%has_field(w_fname) .and. &
     xin%has_field('orog_filt') .and. &
     xin%has_field(ud_fname) .and. xin%has_field(vd_fname) ) then
 
  ! Shortcuts
  atm => self%Atm(1)
  isc = atm%bd%is
  iec = atm%bd%ie
  jsc = atm%bd%js
  jec = atm%bd%je
  npz = atm%npz
 
  ! Orography
  call xin%get_field('orog_filt', orog_filt)
  atm%phis(isc:iec,jsc:jec) = real(orog_filt(isc:iec,jsc:jec,1),kind_fv3)*grav  ! Convert to phis
 
  ! Dynamics fields
  call xin%get_field(ps_fname, ps_cold)
  call xin%get_field(zh_fname, zh_cold)
  call xin%get_field( w_fname,  w_cold)
 
  ! akbk from cold starts have different levels (for now extra levels are zero)
  levp = size(w_cold,3)
  itoa = levp - npz + 1
 
  allocate(ak(levp+1))
  allocate(bk(levp+1))
  ak = 0.0_kind_fv3
  bk = 0.0_kind_fv3
  ak(itoa:levp+1) = atm%ak(1:npz+1)
  bk(itoa:levp+1) = atm%bk(1:npz+1)
  ak(1) = max(1.e-9_kind_fv3, ak(1))
 
  ! Tracers
  ! -------
  ! Number of tracers in Atm
  call get_number_tracers(model_atmos, num_tracers=ntracers, num_prog=ntprog)
  call get_number_tracers(model_atmos, num_tracers=ntracers, num_prog=ntprog)
 
  ! initialize all tracers to default values prior to being input
  do nt = 1, ntprog
    call get_tracer_names(model_atmos, nt, tracer_name)
    ! set all tracers to an initial profile value
    call set_tracer_profile (model_atmos, nt, atm%q(:,:,:,nt)  )
  enddo
  do nt = ntprog+1, ntracers
    call get_tracer_names(model_atmos, nt, tracer_name)
    ! set all tracers to an initial profile value
    call set_tracer_profile (model_atmos, nt, atm%qdiag(:,:,:,nt)  )
  enddo
 
  ! Copy from fv3-jedi to temporary cold q structure
  allocate (q_cold(isc:iec, jsc:jec, levp, ntracers))
  q_cold = 0.0_kind_fv3
  do nt = 1, ntracers
    call get_tracer_names(model_atmos, nt, tracer_name)
    if (self%from_cold_start) tracer_name = trim(tracer_name)//"_cold"
    if (xin%has_field(trim(tracer_name))) then
      call xin%get_field(trim(tracer_name), q_tmp)
      q_cold(:,:,:,nt) = real(q_tmp, kind_fv3)
    endif
  enddo
 
  ! Call remapping non-wind variables
  ! ---------------------------------
  if (xin%has_field(t_fname)) then
    call xin%get_field( t_fname,  t_cold)
    call remap_scalar(atm, levp, npz, ntracers, ak, bk, real(ps_cold(:,:,1),kind_fv3), q_cold, &
                      real(zh_cold,kind_fv3), real(w_cold,kind_fv3), real(t_cold,kind_fv3))
  else
    call remap_scalar(atm, levp, npz, ntracers, ak, bk, real(ps_cold(:,:,1),kind_fv3), q_cold, &
                      real(zh_cold,kind_fv3), real(w_cold,kind_fv3))
  endif
 
  ! Call remapping wind variables
  ! -----------------------------
  call xin%get_field(ud_fname, ud_cold)
  call xin%get_field(vd_fname, vd_cold)
 
  call remap_dwinds(levp, npz, ak, bk, real(ps_cold(:,:,1),kind_fv3), real(ud_cold,kind_fv3), &
                    real(vd_cold,kind_fv3), Atm)
 
  ! Tracer weighting
  ! ----------------
  liq_wat = get_tracer_index(model_atmos, 'liq_wat')
  ice_wat = get_tracer_index(model_atmos, 'ice_wat')
  rainwat = get_tracer_index(model_atmos, 'rainwat')
  snowwat = get_tracer_index(model_atmos, 'snowwat')
  graupel = get_tracer_index(model_atmos, 'graupel')
  ntclamt = get_tracer_index(model_atmos, 'cld_amt')
 
  if (self%from_cold_start) then
    do k = 1,npz
      do j = jsc,jec
        do i = isc,iec
          wt = atm%delp(i,j,k)
          if ( atm%flagstruct%nwat == 6 ) then
            qt = wt*(1.0_kind_fv3 + atm%q(i,j,k,liq_wat) + atm%q(i,j,k,ice_wat) + &
                                    atm%q(i,j,k,rainwat) + atm%q(i,j,k,snowwat) + &
                                    atm%q(i,j,k,graupel))
          else
            qt = wt*(1.0_kind_fv3 + sum(atm%q(i,j,k,2:atm%flagstruct%nwat)))
          endif
          atm%delp(i,j,k) = qt
          if (ntclamt > 0) atm%q(i,j,k,ntclamt) = 0.0
        enddo
      enddo
    enddo
  else
    ! TODO Question is do we need this if just remapping after adding the increment, say?
    do k = 1,npz
      do j = jsc,jec
        do i = isc,iec
          wt = atm%delp(i,j,k)
          if ( atm%flagstruct%nwat == 6 ) then
            qt = wt*(1.0_kind_fv3 + atm%q(i,j,k,liq_wat) + atm%q(i,j,k,ice_wat) + &
                                    atm%q(i,j,k,rainwat) + atm%q(i,j,k,snowwat) + &
                                    atm%q(i,j,k,graupel))
          else
             qt = wt*(1.0_kind_fv3 + sum(atm%q(i,j,k,2:atm%flagstruct%nwat)))
          endif
          m_fac = wt / qt
          do nt=1,ntracers
            atm%q(i,j,k,nt) = m_fac * atm%q(i,j,k,nt)
          enddo
          atm%delp(i,j,k) = qt
          if (ntclamt > 0) atm%q(i,j,k,ntclamt) = 0.0
        enddo
      enddo
    enddo
  endif
 
  if (self%checker_tr) then
    nts = ntracers - self%nt_checker+1
    call checker_tracers(isc, iec, jsc, jec, atm%bd%isd, atm%bd%ied, atm%bd%jsd, atm%bd%jed, &
                         self%nt_checker, npz, atm%q(:,:,:,nts:ntracers), &
                         atm%gridstruct%agrid_64(isc:iec,jsc:jec,1),     &
                         atm%gridstruct%agrid_64(isc:iec,jsc:jec,2), &
                         real(9.0_kind_real,kind_fv3), real(9.0_kind_real,kind_fv3))
  endif
 
  have_remapped = .true.
 
endif
 
 
! Loop over the fields not found in the input state and work through cases
! ------------------------------------------------------------------------
do f = 1, size(fields_to_do)
 
  ! Remapping needs to have happened in order to copy anything
  if (.not. have_remapped) call field_fail(fields_to_do(f))
 
  call xout%get_field(trim(fields_to_do(f)), field_ptr)
 
  select case (trim(fields_to_do(f)))
 
  case ("ud")
 
    field_ptr%array(isc:iec,jsc:jec+1,:) = atm%u(isc:iec,jsc:jec+1,:)
 
  case ("vd")
 
    field_ptr%array(isc:iec+1,jsc:jec,:) = atm%v(isc:iec+1,jsc:jec,:)
 
  case ("t")
 
    field_ptr%array(isc:iec,jsc:jec,:) = atm%pt(isc:iec,jsc:jec,:)
 
  case ("delp")
 
    field_ptr%array(isc:iec,jsc:jec,:) = atm%delp(isc:iec,jsc:jec,:)
 
  case ("ps")
 
    field_ptr%array(isc:iec,jsc:jec,1) = atm%ps(isc:iec,jsc:jec)
 
  case ("delz")
 
    field_ptr%array(isc:iec,jsc:jec,:) = atm%delz(isc:iec,jsc:jec,:)
 
  case ("w")
 
    field_ptr%array(isc:iec,jsc:jec,:) = atm%w(isc:iec,jsc:jec,:)
 
  case ("phis")
 
    field_ptr%array(isc:iec,jsc:jec,1) = atm%phis(isc:iec,jsc:jec)
 
  case ('sphum','liq_wat','ice_wat','o3mr','graupel','snowwat','rainwat')
 
    nt = get_tracer_index(model_atmos, trim(fields_to_do(f)))
    if (nt == no_tracer) call field_fail(fields_to_do(f))
    field_ptr%array(isc:iec,jsc:jec,:) = atm%q(isc:iec,jsc:jec,:,nt)
 
  case ('sgs_tke','cld_amt')
 
    field_ptr%array = 0.0_kind_real
 
  case default
 
    call abor1_ftn("fv3jedi_vc_coldstartwinds_mod.changevar unknown field: "//trim(fields_to_do(f))&
                   //". Not in input field and no transform case specified.")
 
  end select
 
enddo
 
 
! Copy calendar infomation
! ------------------------
xout%calendar_type = xin%calendar_type
xout%date_init = xin%date_init
 
end subroutine changevar
 
! --------------------------------------------------------------------------------------------------
 
end module fv3jedi_vc_vertremap_mod