fv_io.F90

Go to the documentation of this file.

!***********************************************************************
!*                   GNU Lesser General Public License
!*
!* This file is part of the FV3 dynamical core.
!*
!* The FV3 dynamical core is free software: you can redistribute it
!* and/or modify it under the terms of the
!* GNU Lesser General Public License as published by the
!* Free Software Foundation, either version 3 of the License, or
!* (at your option) any later version.
!*
!* The FV3 dynamical core is distributed in the hope that it will be
!* useful, but WITHOUT ANYWARRANTY; without even the implied warranty
!* of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
!* See the GNU General Public License for more details.
!*
!* You should have received a copy of the GNU Lesser General Public
!* License along with the FV3 dynamical core.
!* If not, see <http://www.gnu.org/licenses/>.
!***********************************************************************


module fv_io_mod

! <table>
! <tr>
!     <th>Module Name</th>
!     <th>Functions Included</th>
!   </tr>
!   <tr>
!     <td>external_sst_mod</td>
!     <td>sst_ncep, sst_anom, use_ncep_sst</td>
!   </tr>
!   <tr>
!     <td>field_manager_mod</td>
!     <td>MODEL_ATMOS</td>
!   </tr>
!   <tr>
!     <td>fms_mod</td>
!     <td>file_exist</td>
!   </tr>
!   <tr>
!     <td>fms_io_mod</td>
!     <td>fms_io_exit, get_tile_string,restart_file_type, 
!         register_restart_field, save_restart, restore_state, 
!         set_domain, nullify_domain, set_filename_appendix, 
!         get_mosaic_tile_file, get_instance_filename, 
!         save_restart_border, restore_state_border,
!         free_restart_type,field_exist</td>
!   </tr>
!   <tr>
!     <td>fv_arrays_mod</td>
!     <td>fv_atmos_type, fv_nest_BC_type_3D</td>
!   </tr>
!   <tr>
!     <td>fv_eta_mod</td>
!     <td>set_external_eta</td>
!   </tr>
!   <tr>
!     <td>fv_mp_mod</td>
!     <td>ng, mp_gather, is_master</td>
!   </tr>
!   <tr>
!     <td>mpp_mod</td>
!     <td>mpp_error, FATAL, NOTE, WARNING, mpp_root_pe,
!         mpp_sync, mpp_pe, mpp_declare_pelist</td>
!   </tr>
!   <tr>
!     <td>mpp_domains_mod</td>
!     <td>domain2d, EAST, WEST, NORTH, CENTER, SOUTH, CORNER,
!         mpp_get_compute_domain, mpp_get_data_domain,  
!         mpp_get_layout, mpp_get_ntile_count,mpp_get_global_domain</td>
!   </tr>
!   <tr>
!     <td>tracer_manager_mod</td>
!     <td>tr_get_tracer_names=>get_tracer_names, 
!         get_tracer_names, get_number_tracers, 
!         set_tracer_profile, get_tracer_index</td>
!   </tr>
! </table>

  use fms_mod,                 only: file_exist
  use fms_io_mod,              only: fms_io_exit, get_tile_string, &
                                     restart_file_type, register_restart_field, &
                                     save_restart, restore_state, &
                                     set_domain, nullify_domain, set_filename_appendix, &
                                     get_mosaic_tile_file, get_instance_filename, & 
                                     save_restart_border, restore_state_border, free_restart_type, &
                                     field_exist
  use mpp_mod,                 only: mpp_error, fatal, note, warning, mpp_root_pe, &
                                     mpp_sync, mpp_pe, mpp_declare_pelist
  use mpp_domains_mod,         only: domain2d, east, west, north, center, south, corner, &
                                     mpp_get_compute_domain, mpp_get_data_domain, & 
                                     mpp_get_layout, mpp_get_ntile_count, &
                                     mpp_get_global_domain
  use tracer_manager_mod,      only: tr_get_tracer_names=>get_tracer_names, &
                                     get_tracer_names, get_number_tracers, &
                                     set_tracer_profile, &
                                     get_tracer_index
  use field_manager_mod,       only: model_atmos  
  use external_sst_mod,        only: sst_ncep, sst_anom, use_ncep_sst
  use fv_arrays_mod,           only: fv_atmos_type, fv_nest_bc_type_3d
  use fv_eta_mod,              only: set_external_eta

  use fv_mp_mod,               only: mp_gather, is_master
  use fms_io_mod,              only: set_domain
  use fv_treat_da_inc_mod,     only: read_da_inc

  implicit none
  private

  public :: fv_io_init, fv_io_exit, fv_io_read_restart, remap_restart, fv_io_write_restart
  public :: fv_io_read_tracers, fv_io_register_restart, fv_io_register_nudge_restart
  public :: fv_io_register_restart_bcs
  public :: fv_io_write_bcs, fv_io_read_bcs

  logical                       :: module_is_initialized = .false.


  integer ::grid_xtdimid, grid_ytdimid, haloid, pfullid !For writing BCs
  integer ::grid_xtstagdimid, grid_ytstagdimid, oneid

contains 


  subroutine fv_io_init()
    module_is_initialized = .true.
  end subroutine fv_io_init


  subroutine fv_io_exit
    module_is_initialized = .false.
  end subroutine fv_io_exit


  subroutine  fv_io_read_restart(fv_domain,Atm)
    type(domain2d),      intent(inout) :: fv_domain
    type(fv_atmos_type), intent(inout) :: Atm(:)

    character(len=64)    :: fname, tracer_name
    character(len=6)  :: stile_name
    integer              :: isc, iec, jsc, jec, n, nt, nk, ntracers
    integer              :: ntileMe
    integer              :: ks, ntiles
    real                 :: ptop

    character(len=128)           :: tracer_longname, tracer_units

    ntileme = size(atm(:))  ! This will need mods for more than 1 tile per pe

    call restore_state(atm(1)%Fv_restart)
    if (atm(1)%flagstruct%external_eta) then
       call set_external_eta(atm(1)%ak, atm(1)%bk, atm(1)%ptop, atm(1)%ks)
    endif

    if ( use_ncep_sst .or. atm(1)%flagstruct%nudge .or. atm(1)%flagstruct%ncep_ic ) then
       call mpp_error(note, 'READING FROM SST_RESTART DISABLED')
       !call restore_state(Atm(1)%SST_restart)
    endif

! fix for single tile runs where you need fv_core.res.nc and fv_core.res.tile1.nc
    ntiles = mpp_get_ntile_count(fv_domain)
    if(ntiles == 1 .and. .not. atm(1)%neststruct%nested) then
       stile_name = '.tile1'
    else
       stile_name = ''
    endif
 
    do n = 1, ntileme
       call restore_state(atm(n)%Fv_tile_restart)

!--- restore data for fv_tracer - if it exists
       fname = 'INPUT/fv_tracer.res'//trim(stile_name)//'.nc'
       if (file_exist(fname)) then
         call restore_state(atm(n)%Tra_restart)
       else
         call mpp_error(note,'==> Warning from fv_read_restart: Expected file '//trim(fname)//' does not exist')
       endif

!--- restore data for surface winds - if it exists
       fname = 'INPUT/fv_srf_wnd.res'//trim(stile_name)//'.nc'
       if (file_exist(fname)) then
         call restore_state(atm(n)%Rsf_restart)
         atm(n)%flagstruct%srf_init = .true.
       else
         call mpp_error(note,'==> Warning from fv_read_restart: Expected file '//trim(fname)//' does not exist')
         atm(n)%flagstruct%srf_init = .false.
       endif

       if ( atm(n)%flagstruct%fv_land ) then
!--- restore data for mg_drag - if it exists
         fname = 'INPUT/mg_drag.res'//trim(stile_name)//'.nc'
         if (file_exist(fname)) then
           call restore_state(atm(n)%Mg_restart)
         else
           call mpp_error(note,'==> Warning from fv_read_restart: Expected file '//trim(fname)//' does not exist')
         endif
!--- restore data for fv_land - if it exists
         fname = 'INPUT/fv_land.res'//trim(stile_name)//'.nc'
         if (file_exist(fname)) then
           call restore_state(atm(n)%Lnd_restart)
         else
           call mpp_error(note,'==> Warning from fv_read_restart: Expected file '//trim(fname)//' does not exist')
         endif
       endif

    end do

    return

  end subroutine  fv_io_read_restart
  
  subroutine fv_io_read_tracers(fv_domain,Atm)
    type(domain2d),      intent(inout) :: fv_domain
    type(fv_atmos_type), intent(inout) :: Atm(:)
    integer :: n, ntracers, ntprog, nt, isc, iec, jsc, jec, id_restart
    character(len=6) :: stile_name
    character(len=64):: fname, tracer_name
    type(restart_file_type) :: Tra_restart_r
    integer :: ntiles

    n = 1
    isc = atm(n)%bd%isc
    iec = atm(n)%bd%iec
    jsc = atm(n)%bd%jsc
    jec = atm(n)%bd%jec
    call get_number_tracers(model_atmos, num_tracers=ntracers, num_prog=ntprog)

! fix for single tile runs where you need fv_core.res.nc and fv_core.res.tile1.nc
    ntiles = mpp_get_ntile_count(fv_domain)
    if(ntiles == 1 .and. .not. atm(1)%neststruct%nested) then
       stile_name = '.tile1'
    else
       stile_name = ''
    endif

    fname = 'fv_tracer.res'//trim(stile_name)//'.nc'
    do nt = 2, ntprog
       call get_tracer_names(model_atmos, nt, tracer_name)
       call set_tracer_profile (model_atmos, nt, atm(n)%q(isc:iec,jsc:jec,:,nt)  )
       id_restart = register_restart_field(tra_restart_r, fname, tracer_name, atm(n)%q(:,:,:,nt), &
                    domain=fv_domain, mandatory=.false., tile_count=n)
    enddo
    do nt = ntprog+1, ntracers
       call get_tracer_names(model_atmos, nt, tracer_name)
       call set_tracer_profile (model_atmos, nt, atm(n)%qdiag(isc:iec,jsc:jec,:,nt)  )
       id_restart = register_restart_field(tra_restart_r, fname, tracer_name, atm(n)%qdiag(:,:,:,nt), &
                    domain=fv_domain, mandatory=.false., tile_count=n)
    enddo
    if (file_exist('INPUT'//trim(fname))) then
      call restore_state(tra_restart_r)
      call free_restart_type(tra_restart_r)
    else
      call mpp_error(note,'==> Warning from fv_io_read_tracers: Expected file '//trim(fname)//' does not exist')
    endif

    return

  end subroutine  fv_io_read_tracers
 
  subroutine  remap_restart(fv_domain,Atm)
  use fv_mapz_mod,       only: rst_remap

    type(domain2d),      intent(inout) :: fv_domain
    type(fv_atmos_type), intent(inout) :: Atm(:)

    character(len=64)    :: fname, tracer_name
    character(len=6)     :: stile_name
    integer              :: isc, iec, jsc, jec, n, nt, nk, ntracers, ntprog, ntdiag
    integer              :: isd, ied, jsd, jed
    integer              :: ntiles
    type(restart_file_type) :: FV_restart_r, FV_tile_restart_r, Tra_restart_r
    integer :: id_restart

!
!-------------------------------------------------------------------------
    real, allocatable:: ak_r(:), bk_r(:)
    real, allocatable:: u_r(:,:,:), v_r(:,:,:), pt_r(:,:,:), delp_r(:,:,:)
    real, allocatable:: w_r(:,:,:), delz_r(:,:,:), ze0_r(:,:,:)
    real, allocatable:: q_r(:,:,:,:), qdiag_r(:,:,:,:)
!-------------------------------------------------------------------------
    integer npz, npz_rst, ng
    integer i,j,k

    npz     = atm(1)%npz       ! run time z dimension
    npz_rst = atm(1)%flagstruct%npz_rst   ! restart z dimension
    isc = atm(1)%bd%isc; iec = atm(1)%bd%iec; jsc = atm(1)%bd%jsc; jec = atm(1)%bd%jec
    ng = atm(1)%ng

    isd = isc - ng;  ied = iec + ng
    jsd = jsc - ng;  jed = jec + ng


!   call get_number_tracers(MODEL_ATMOS, num_tracers=ntracers)
    ntprog = size(atm(1)%q,4)  ! Temporary until we get tracer manager integrated
    ntdiag = size(atm(1)%qdiag,4)
    ntracers = ntprog+ntdiag

!    ntileMe = size(Atm(:))  ! This will have to be modified for mult tiles per PE


! Allocate arrays for reading old restart file:
    allocate ( ak_r(npz_rst+1) )
    allocate ( bk_r(npz_rst+1) )

    allocate ( u_r(isc:iec,  jsc:jec+1,npz_rst) )
    allocate ( v_r(isc:iec+1,jsc:jec  ,npz_rst) )

    allocate (   pt_r(isc:iec, jsc:jec,  npz_rst) )
    allocate ( delp_r(isc:iec, jsc:jec,  npz_rst) )
    allocate (    q_r(isc:iec, jsc:jec,  npz_rst, ntprog) )
    allocate (qdiag_r(isc:iec, jsc:jec,  npz_rst, ntprog+1:ntracers) )

    if ( (.not.atm(1)%flagstruct%hydrostatic) .and. (.not.atm(1)%flagstruct%make_nh) ) then
           allocate (    w_r(isc:iec, jsc:jec,  npz_rst) )
           allocate ( delz_r(isc:iec, jsc:jec,  npz_rst) )
           if ( atm(1)%flagstruct%hybrid_z )   &
           allocate ( ze0_r(isc:iec, jsc:jec,  npz_rst+1) )
    endif

    fname = 'fv_core.res.nc'
    id_restart = register_restart_field(fv_restart_r, fname, 'ak', ak_r(:), no_domain=.true.)
    id_restart = register_restart_field(fv_restart_r, fname, 'bk', bk_r(:), no_domain=.true.)
    call restore_state(fv_restart_r)
    call free_restart_type(fv_restart_r)

! fix for single tile runs where you need fv_core.res.nc and fv_core.res.tile1.nc
    ntiles = mpp_get_ntile_count(fv_domain)
    if(ntiles == 1 .and. .not. atm(1)%neststruct%nested) then
       stile_name = '.tile1'
    else
       stile_name = ''
    endif

!!!! A NOTE about file names
!!! file_exist() needs the full relative path, including INPUT/
!!! But register_restart_field ONLY looks in INPUT/ and so JUST needs the file name!!

!    do n = 1, ntileMe
    n = 1
       fname = 'fv_core.res'//trim(stile_name)//'.nc'
       id_restart =  register_restart_field(fv_tile_restart_r, fname, 'u', u_r, &
                     domain=fv_domain, position=north,tile_count=n)
       id_restart =  register_restart_field(fv_tile_restart_r, fname, 'v', v_r, &
                     domain=fv_domain, position=east,tile_count=n)
       if (.not.atm(n)%flagstruct%hydrostatic) then
          id_restart =  register_restart_field(fv_tile_restart_r, fname, 'W', w_r, &
                        domain=fv_domain, mandatory=.false., tile_count=n)
          id_restart =  register_restart_field(fv_tile_restart_r, fname, 'DZ', delz_r, &
                        domain=fv_domain, mandatory=.false., tile_count=n)
          if ( atm(n)%flagstruct%hybrid_z ) then
             id_restart =  register_restart_field(fv_tile_restart_r, fname, 'ZE0', ze0_r, &
                           domain=fv_domain, mandatory=.false., tile_count=n)
          endif
       endif
       id_restart =  register_restart_field(fv_tile_restart_r, fname, 'T', pt_r, &
                     domain=fv_domain, tile_count=n)
       id_restart =  register_restart_field(fv_tile_restart_r, fname, 'delp', delp_r, &
                     domain=fv_domain, tile_count=n)
       id_restart =  register_restart_field(fv_tile_restart_r, fname, 'phis', atm(n)%phis, &
                     domain=fv_domain, tile_count=n)
       call restore_state(fv_tile_restart_r)
       call free_restart_type(fv_tile_restart_r)
       fname = 'fv_srf_wnd.res'//trim(stile_name)//'.nc'
       if (file_exist('INPUT/'//fname)) then
         call restore_state(atm(n)%Rsf_restart)
         atm(n)%flagstruct%srf_init = .true.
       else
         call mpp_error(note,'==> Warning from remap_restart: Expected file '//trim(fname)//' does not exist')
         atm(n)%flagstruct%srf_init = .false.
       endif

       if ( atm(n)%flagstruct%fv_land ) then
!--- restore data for mg_drag - if it exists
         fname = 'mg_drag.res'//trim(stile_name)//'.nc'
         if (file_exist('INPUT/'//fname)) then
           call restore_state(atm(n)%Mg_restart)
         else
           call mpp_error(note,'==> Warning from remap_restart: Expected file '//trim(fname)//' does not exist')
         endif
!--- restore data for fv_land - if it exists
         fname = 'fv_land.res'//trim(stile_name)//'.nc'
         if (file_exist('INPUT/'//fname)) then
           call restore_state(atm(n)%Lnd_restart)
         else
           call mpp_error(note,'==> Warning from remap_restart: Expected file '//trim(fname)//' does not exist')
         endif
       endif

       fname = 'fv_tracer.res'//trim(stile_name)//'.nc'
       if (file_exist('INPUT/'//fname)) then
         do nt = 1, ntprog
            call get_tracer_names(model_atmos, nt, tracer_name)
            call set_tracer_profile (model_atmos, nt, q_r(isc:iec,jsc:jec,:,nt)  )
            id_restart = register_restart_field(tra_restart_r, fname, tracer_name, q_r(:,:,:,nt), &
                         domain=fv_domain, mandatory=.false., tile_count=n)
         enddo
         do nt = ntprog+1, ntracers
            call get_tracer_names(model_atmos, nt, tracer_name)
            call set_tracer_profile (model_atmos, nt, qdiag_r(isc:iec,jsc:jec,:,nt)  )
            id_restart = register_restart_field(tra_restart_r, fname, tracer_name, qdiag_r(:,:,:,nt), &
                         domain=fv_domain, mandatory=.false., tile_count=n)
         enddo
         call restore_state(tra_restart_r)
         call free_restart_type(tra_restart_r)
       else
         call mpp_error(note,'==> Warning from remap_restart: Expected file '//trim(fname)//' does not exist')
       endif

!      ====== PJP added DA functionality ======
       if (atm(n)%flagstruct%read_increment) then
          ! print point in middle of domain for a sanity check
          i = (isc + iec)/2
          j = (jsc + jec)/2
          k = npz_rst/2
          if( is_master() ) write(*,*) 'Calling read_da_inc',pt_r(i,j,k)
          call read_da_inc(atm(n), atm(n)%domain, atm(n)%bd, npz_rst, ntprog, &
               u_r, v_r, q_r, delp_r, pt_r, delz_r, isc, jsc, iec, jec, &
                isc, jsc, iec, jec )
          if( is_master() ) write(*,*) 'Back from read_da_inc',pt_r(i,j,k)
       endif
!      ====== end PJP added DA functionailty======
       
       call rst_remap(npz_rst, npz, isc, iec, jsc, jec, isd, ied, jsd, jed, ntracers, ntprog,      &
                      delp_r,      u_r,      v_r,      w_r,      delz_r,      pt_r,  q_r,  qdiag_r,&
                      atm(n)%delp, atm(n)%u, atm(n)%v, atm(n)%w, atm(n)%delz, atm(n)%pt, atm(n)%q, &
                      atm(n)%qdiag, ak_r,  bk_r, atm(n)%ptop, atm(n)%ak, atm(n)%bk,                &
                      atm(n)%flagstruct%hydrostatic, atm(n)%flagstruct%make_nh, atm(n)%domain,     &
                      atm(n)%gridstruct%square_domain)
    !end do

    deallocate( ak_r )
    deallocate( bk_r )
    deallocate( u_r )
    deallocate( v_r )
    deallocate( pt_r )
    deallocate( delp_r )
    deallocate( q_r )
    deallocate( qdiag_r )

    if ( (.not.atm(1)%flagstruct%hydrostatic) .and. (.not.atm(1)%flagstruct%make_nh) ) then
         deallocate ( w_r )
         deallocate ( delz_r )
         if ( atm(1)%flagstruct%hybrid_z ) deallocate ( ze0_r )
    endif

  end subroutine  remap_restart

  subroutine  fv_io_register_nudge_restart(Atm)
    type(fv_atmos_type), intent(inout) :: Atm(:)
    character(len=64) :: fname
    integer           :: id_restart

! use_ncep_sst may not be initialized at this point?
    call mpp_error(note, 'READING FROM SST_restart DISABLED')
    if ( use_ncep_sst .or. atm(1)%flagstruct%nudge .or. atm(1)%flagstruct%ncep_ic ) then
!    if ( Atm(1)%nudge .or. Atm(1)%ncep_ic ) then
       fname = 'sst_ncep.res.nc'
       id_restart = register_restart_field(atm(1)%SST_restart, fname, 'sst_ncep', sst_ncep)
       id_restart = register_restart_field(atm(1)%SST_restart, fname, 'sst_anom', sst_anom)
    endif

  end subroutine  fv_io_register_nudge_restart

  subroutine  fv_io_register_restart(fv_domain,Atm)
    type(domain2d),      intent(inout) :: fv_domain
    type(fv_atmos_type), intent(inout) :: Atm(:)

    character(len=64) :: fname, tracer_name
    character(len=6)  :: gn, stile_name
    integer           :: id_restart
    integer           :: n, nt, ntracers, ntprog, ntdiag, ntileMe, ntiles

    ntileme = size(atm(:)) 
    ntprog = size(atm(1)%q,4) 
    ntdiag = size(atm(1)%qdiag,4) 
    ntracers = ntprog+ntdiag

!--- set the 'nestXX' appendix for all files using fms_io
    if (atm(1)%grid_number > 1) then
       write(gn,'(A4, I2.2)') "nest", atm(1)%grid_number
    else
       gn = ''
    end if
    call set_filename_appendix(gn)

!--- fix for single tile runs where you need fv_core.res.nc and fv_core.res.tile1.nc
    ntiles = mpp_get_ntile_count(fv_domain)
    if(ntiles == 1 .and. .not. atm(1)%neststruct%nested) then
       stile_name = '.tile1'
    else
       stile_name = ''
    endif

! use_ncep_sst may not be initialized at this point?
#ifndef DYCORE_SOLO
!    call mpp_error(NOTE, 'READING FROM SST_RESTART DISABLED')
!!$   if ( use_ncep_sst .or. Atm(1)%flagstruct%nudge .or. Atm(1)%flagstruct%ncep_ic ) then
!!$       fname = 'sst_ncep'//trim(gn)//'.res.nc'
!!$       id_restart = register_restart_field(Atm(1)%SST_restart, fname, 'sst_ncep', sst_ncep)
!!$       id_restart = register_restart_field(Atm(1)%SST_restart, fname, 'sst_anom', sst_anom)
!!$   endif
#endif

    fname = 'fv_core.res.nc'
    id_restart = register_restart_field(atm(1)%Fv_restart, fname, 'ak', atm(1)%ak(:), no_domain=.true.)
    id_restart = register_restart_field(atm(1)%Fv_restart, fname, 'bk', atm(1)%bk(:), no_domain=.true.) 

    do n = 1, ntileme
       fname = 'fv_core.res'//trim(stile_name)//'.nc'
       id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'u', atm(n)%u, &
                     domain=fv_domain, position=north,tile_count=n)
       id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'v', atm(n)%v, &
                     domain=fv_domain, position=east,tile_count=n)
       if (.not.atm(n)%flagstruct%hydrostatic) then
          id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'W', atm(n)%w, &
                        domain=fv_domain, mandatory=.false., tile_count=n)
          id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'DZ', atm(n)%delz, &
                        domain=fv_domain, mandatory=.false., tile_count=n)
          if ( atm(n)%flagstruct%hybrid_z ) then
             id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'ZE0', atm(n)%ze0, &
                           domain=fv_domain, mandatory=.false., tile_count=n)
          endif
       endif
       id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'T', atm(n)%pt, &
                     domain=fv_domain, tile_count=n)
       id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'delp', atm(n)%delp, &
                     domain=fv_domain, tile_count=n)
       id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'phis', atm(n)%phis, &
                     domain=fv_domain, tile_count=n)

       !--- include agrid winds in restarts for use in data assimilation 
       if (atm(n)%flagstruct%agrid_vel_rst) then
         id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'ua', atm(n)%ua, &
                       domain=fv_domain, tile_count=n, mandatory=.false.)
         id_restart =  register_restart_field(atm(n)%Fv_tile_restart, fname, 'va', atm(n)%va, &
                       domain=fv_domain, tile_count=n, mandatory=.false.)
       endif

       fname = 'fv_srf_wnd.res'//trim(stile_name)//'.nc'
       id_restart =  register_restart_field(atm(n)%Rsf_restart, fname, 'u_srf', atm(n)%u_srf, &
                     domain=fv_domain, tile_count=n)
       id_restart =  register_restart_field(atm(n)%Rsf_restart, fname, 'v_srf', atm(n)%v_srf, &
                     domain=fv_domain, tile_count=n)
#ifdef SIM_PHYS
       id_restart =  register_restart_field(rsf_restart(n), fname, 'ts', atm(n)%ts, &
                     domain=fv_domain, tile_count=n)
#endif

       if ( atm(n)%flagstruct%fv_land ) then
          !-------------------------------------------------------------------------------------------------
          ! Optional terrain deviation (sgh) and land fraction (oro)
          fname = 'mg_drag.res'//trim(stile_name)//'.nc'
          id_restart =  register_restart_field(atm(n)%Mg_restart, fname, 'ghprime', atm(n)%sgh, &
                        domain=fv_domain, tile_count=n)  

          fname = 'fv_land.res'//trim(stile_name)//'.nc'
          id_restart = register_restart_field(atm(n)%Lnd_restart, fname, 'oro', atm(n)%oro, &
                        domain=fv_domain, tile_count=n)
       endif

       fname = 'fv_tracer.res'//trim(stile_name)//'.nc'
       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(n)%q(:,:,:,nt)  )
          id_restart = register_restart_field(atm(n)%Tra_restart, fname, tracer_name, atm(n)%q(:,:,:,nt), &
                       domain=fv_domain, mandatory=.false., tile_count=n)
       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(n)%qdiag(:,:,:,nt)  )
          id_restart = register_restart_field(atm(n)%Tra_restart, fname, tracer_name, atm(n)%qdiag(:,:,:,nt), &
                       domain=fv_domain, mandatory=.false., tile_count=n)
       enddo

       if ( atm(n)%neststruct%nested ) then
          call fv_io_register_restart_bcs(atm(n)) !TODO put into fv_io_register_restart
       endif

    enddo

  end subroutine  fv_io_register_restart

  subroutine  fv_io_write_restart(Atm, timestamp)
    type(fv_atmos_type),        intent(inout) :: Atm
    character(len=*), optional, intent(in) :: timestamp

!!$    if ( use_ncep_sst .or. Atm%flagstruct%nudge .or. Atm%flagstruct%ncep_ic ) then
!!$       call mpp_error(NOTE, 'READING FROM SST_RESTART DISABLED')
!!$       !call save_restart(Atm%SST_restart, timestamp)
!!$    endif
 
    if ( (use_ncep_sst .or. atm%flagstruct%nudge) .and. .not. atm%gridstruct%nested ) then
       call save_restart(atm%SST_restart, timestamp)
    endif

    call save_restart(atm%Fv_restart, timestamp)
    call save_restart(atm%Fv_tile_restart, timestamp)
    call save_restart(atm%Rsf_restart, timestamp)

    if ( atm%flagstruct%fv_land ) then
       call save_restart(atm%Mg_restart, timestamp)
       call save_restart(atm%Lnd_restart, timestamp)
    endif

    call save_restart(atm%Tra_restart, timestamp)


  end subroutine  fv_io_write_restart

  subroutine register_bcs_2d(Atm, BCfile_ne, BCfile_sw, fname_ne, fname_sw, &
                             var_name, var, var_bc, istag, jstag)
    type(fv_atmos_type),      intent(in)    :: Atm
    type(restart_file_type),  intent(inout) :: BCfile_ne, BCfile_sw
    character(len=120),       intent(in)    :: fname_ne, fname_sw
    character(len=*),         intent(in)    :: var_name
    real, dimension(:,:),     intent(in), optional :: var
    type(fv_nest_bc_type_3d), intent(in), optional :: var_bc
    integer,                  intent(in), optional :: istag, jstag

    integer :: npx, npy, i_stag, j_stag
    integer :: is, ie, js, je, isd, ied, jsd, jed, n
    integer :: x_halo, y_halo, x_halo_ns, id_restart
    integer :: layout(2), global_size(2), indices(4)
    integer, allocatable, dimension(:) :: x1_pelist, y1_pelist
    integer, allocatable, dimension(:) :: x2_pelist, y2_pelist
    logical :: is_root_pe

    i_stag = 0 
    j_stag = 0 
    if (present(istag)) i_stag = i_stag
    if (present(jstag)) j_stag = j_stag
    call mpp_get_global_domain(atm%domain, xsize = npx, ysize = npy, position=corner )
    call mpp_get_data_domain(atm%domain, isd, ied, jsd, jed )
    call mpp_get_compute_domain(atm%domain, is, ie, js, je )
    call mpp_get_layout(atm%domain, layout)
    allocate (x1_pelist(layout(1)))
    allocate (y1_pelist(layout(2)))
    allocate (x2_pelist(layout(1)))
    allocate (y2_pelist(layout(2)))
    x_halo = is-isd
    y_halo = js-jsd
! define west and east pelist
    do n = 1,layout(2)
      y1_pelist(n)=mpp_root_pe()+layout(1)*n-1
      y2_pelist(n)=mpp_root_pe()+layout(1)*(n-1)
    enddo
! define south and north pelist
    do n = 1,layout(1)
      x1_pelist(n)=mpp_root_pe()+layout(1)*(layout(2)-1)+(n-1)
      x2_pelist(n)=mpp_root_pe()+(n-1)
    enddo
! declare the pelists inside of mpp (creates the MPI communicator)
    call mpp_declare_pelist(x1_pelist)
    call mpp_declare_pelist(x2_pelist)
    call mpp_declare_pelist(y1_pelist)
    call mpp_declare_pelist(y2_pelist)

!EAST & WEST
!set defaults for west/east halo regions
    indices(1) = 1
    indices(2) = x_halo
    indices(3) = jsd
    indices(4) = jed+j_stag
    global_size(1) = x_halo
    global_size(2) = npy-1+2*y_halo+j_stag

!define west root_pe
    is_root_pe = .false.
    if (is.eq.1 .and. js.eq.1) is_root_pe = .true.
!register west halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_west_t1', &
                                        var_bc%west_t1, & 
                                        indices, global_size, y2_pelist, &
                                        is_root_pe, jshift=y_halo)
!register west prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_west', &
                                        var, indices, global_size, &
                                        y2_pelist, is_root_pe, jshift=y_halo)

!define east root_pe
    is_root_pe = .false.
    if (ie.eq.npx-1 .and. je.eq.npy-1) is_root_pe = .true.
!register east halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_east_t1', &
                                        var_bc%east_t1, & 
                                        indices, global_size, y1_pelist, &
                                        is_root_pe, jshift=y_halo)

!reset indices for prognostic variables in the east halo
    indices(1) = ied-x_halo+1+i_stag
    indices(2) = ied+i_stag
!register east prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_east', &
                                        var, indices, global_size, &
                                        y1_pelist, is_root_pe, jshift=y_halo, &
                                        x_halo=(size(var,1)-x_halo), ishift=-(ie+i_stag))

!NORTH & SOUTH
!set defaults for north/south halo regions
    indices(1) = isd
    indices(2) = ied+i_stag
    indices(3) = 1
    indices(4) = y_halo
    global_size(1) = npx-1+i_stag
    global_size(2) = y_halo
!modify starts and ends for certain pes
    if (is.eq.1)     indices(1) = is
    if (ie.eq.npx-1) indices(2) = ie+i_stag
    x_halo_ns = 0
    if (is.eq.1) x_halo_ns=x_halo

!define south root_pe
    is_root_pe = .false.
    if (is.eq.1 .and. js.eq.1) is_root_pe = .true.
!register south halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_south_t1', &
                                        var_bc%south_t1, & 
                                        indices, global_size, x2_pelist, &
                                        is_root_pe, x_halo=x_halo_ns)
!register south prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_south', &
                                        var, indices, global_size, &
                                        x2_pelist, is_root_pe, x_halo=x_halo_ns)

!define north root_pe
    is_root_pe = .false.
    if (ie.eq.npx-1 .and. je.eq.npy-1) is_root_pe = .true.
!register north halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_north_t1', &
                                        var_bc%north_t1, & 
                                        indices, global_size, x1_pelist, &
                                        is_root_pe, x_halo=x_halo_ns)

!reset indices for prognostic variables in the north halo
    indices(3) = jed-y_halo+1+j_stag
    indices(4) = jed+j_stag
!register north prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_north', &
                                        var, indices, global_size, &
                                        x1_pelist, is_root_pe, x_halo=x_halo_ns, &
                                        y_halo=(size(var,2)-y_halo), jshift=-(je+j_stag))

  end subroutine register_bcs_2d


  subroutine register_bcs_3d(Atm, BCfile_ne, BCfile_sw, fname_ne, fname_sw, &
                             var_name, var, var_bc, istag, jstag, mandatory)
    type(fv_atmos_type),      intent(in)    :: Atm
    type(restart_file_type),  intent(inout) :: BCfile_ne, BCfile_sw
    character(len=120),       intent(in)    :: fname_ne, fname_sw
    character(len=*),         intent(in)    :: var_name
    real, dimension(:,:,:),   intent(in), optional :: var
    type(fv_nest_bc_type_3d), intent(in), optional :: var_bc
    integer,                  intent(in), optional :: istag, jstag
    logical,                  intent(IN), optional :: mandatory

    integer :: npx, npy, i_stag, j_stag
    integer :: is, ie, js, je, isd, ied, jsd, jed, n
    integer :: x_halo, y_halo, x_halo_ns, id_restart
    integer :: layout(2), global_size(3), indices(4)
    integer, allocatable, dimension(:) :: x1_pelist, y1_pelist
    integer, allocatable, dimension(:) :: x2_pelist, y2_pelist
    logical :: is_root_pe

    i_stag = 0
    j_stag = 0
    if (present(istag)) i_stag = istag
    if (present(jstag)) j_stag = jstag
    call mpp_get_global_domain(atm%domain, xsize = npx, ysize = npy, position=corner )
    call mpp_get_data_domain(atm%domain, isd, ied, jsd, jed )
    call mpp_get_compute_domain(atm%domain, is, ie, js, je )
    call mpp_get_layout(atm%domain, layout)
    allocate (x1_pelist(layout(1)))
    allocate (y1_pelist(layout(2)))
    allocate (x2_pelist(layout(1)))
    allocate (y2_pelist(layout(2)))
    x_halo = is-isd
    y_halo = js-jsd
! define west and east pelist
    do n = 1,layout(2)
      y1_pelist(n)=mpp_root_pe()+layout(1)*n-1
      y2_pelist(n)=mpp_root_pe()+layout(1)*(n-1)
    enddo
! define south and north pelist
    do n = 1,layout(1)
      x1_pelist(n)=mpp_root_pe()+layout(1)*(layout(2)-1)+(n-1)
      x2_pelist(n)=mpp_root_pe()+(n-1)
    enddo
! declare the pelists inside of mpp (creates the MPI communicator)
    call mpp_declare_pelist(x1_pelist)
    call mpp_declare_pelist(x2_pelist)
    call mpp_declare_pelist(y1_pelist)
    call mpp_declare_pelist(y2_pelist)

!EAST & WEST
!set defaults for west/east halo regions
    indices(1) = 1
    indices(2) = x_halo
    indices(3) = jsd
    indices(4) = jed + j_stag
    global_size(1) = x_halo
    global_size(2) = npy-1+2*y_halo + j_stag
    global_size(3) = atm%npz

!define west root_pe
    is_root_pe = .false.
    if (is.eq.1 .and. js.eq.1) is_root_pe = .true.
!register west halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_west_t1', &
                                        var_bc%west_t1, & 
                                        indices, global_size, y2_pelist, &
                                        is_root_pe, jshift=y_halo, mandatory=mandatory)
!register west prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_west', &
                                        var, indices, global_size, &
                                        y2_pelist, is_root_pe, jshift=y_halo, mandatory=mandatory)

!define east root_pe
    is_root_pe = .false.
    if (ie.eq.npx-1 .and. je.eq.npy-1) is_root_pe = .true.
!register east halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_east_t1', &
                                        var_bc%east_t1, & 
                                        indices, global_size, y1_pelist, &
                                        is_root_pe, jshift=y_halo, mandatory=mandatory)

!reset indices for prognostic variables in the east halo
    indices(1) = ied-x_halo+1+i_stag
    indices(2) = ied+i_stag
!register east prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_east', &
                                        var, indices, global_size, &
                                        y1_pelist, is_root_pe, jshift=y_halo, &
                                        x_halo=(size(var,1)-x_halo), ishift=-(ie+i_stag), mandatory=mandatory)

!NORTH & SOUTH
!set defaults for north/south halo regions
    indices(1) = isd
    indices(2) = ied+i_stag
    indices(3) = 1
    indices(4) = y_halo
    global_size(1) = npx-1+i_stag
    global_size(2) = y_halo
    global_size(3) = atm%npz
!modify starts and ends for certain pes
    if (is.eq.1)     indices(1) = is
    if (ie.eq.npx-1) indices(2) = ie+i_stag
    x_halo_ns = 0
    if (is.eq.1) x_halo_ns=x_halo

!define south root_pe
    is_root_pe = .false.
    if (is.eq.1 .and. js.eq.1) is_root_pe = .true.
!register south halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_south_t1', &
                                        var_bc%south_t1, & 
                                        indices, global_size, x2_pelist, &
                                        is_root_pe, x_halo=x_halo_ns, mandatory=mandatory)
!register south prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_sw, trim(fname_sw), &
                                        trim(var_name)//'_south', &
                                        var, indices, global_size, &
                                        x2_pelist, is_root_pe, x_halo=x_halo_ns, mandatory=mandatory)

!define north root_pe
    is_root_pe = .false.
    if (ie.eq.npx-1 .and. je.eq.npy-1) is_root_pe = .true.
!register north halo data in t1
    if (present(var_bc)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_north_t1', &
                                        var_bc%north_t1, & 
                                        indices, global_size, x1_pelist, &
                                        is_root_pe, x_halo=x_halo_ns, mandatory=mandatory)

!reset indices for prognostic variables in the north halo
    indices(3) = jed-y_halo+1+j_stag
    indices(4) = jed+j_stag
!register north prognostic halo data
    if (present(var)) id_restart = register_restart_field(bcfile_ne, trim(fname_ne), &
                                        trim(var_name)//'_north', &
                                        var, indices, global_size, &
                                        x1_pelist, is_root_pe, x_halo=x_halo_ns, &
                                        y_halo=(size(var,2)-y_halo), jshift=-(je+j_stag), mandatory=mandatory)

  end subroutine register_bcs_3d

  subroutine fv_io_register_restart_bcs(Atm)
    type(fv_atmos_type),        intent(inout) :: Atm

    integer :: n, ntracers, ntprog, ntdiag
    character(len=120) :: tname, fname_ne, fname_sw

    fname_ne = 'fv_BC_ne.res.nc'
    fname_sw = 'fv_BC_sw.res.nc'

    ntprog=size(atm%q,4)
    ntdiag=size(atm%qdiag,4)
    ntracers=ntprog+ntdiag

    call set_domain(atm%domain)

    call register_bcs_2d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'phis', var=atm%phis)
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'delp', atm%delp, atm%neststruct%delp_BC)
    do n=1,ntprog
       call get_tracer_names(model_atmos, n, tname)
       call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                            fname_ne, fname_sw, trim(tname), atm%q(:,:,:,n), atm%neststruct%q_BC(n), mandatory=.false.)
    enddo
    do n=ntprog+1,ntracers
       call get_tracer_names(model_atmos, n, tname)
       call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                            fname_ne, fname_sw, trim(tname), var=atm%qdiag(:,:,:,n), mandatory=.false.)
    enddo
#ifndef SW_DYNAMICS
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'pt', atm%pt, atm%neststruct%pt_BC)
    if ((.not.atm%flagstruct%hydrostatic)) then 
       if (is_master()) print*, 'fv_io_register_restart_BCs: REGISTERING NH BCs'
      call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                           fname_ne, fname_sw, 'w', atm%w, atm%neststruct%w_BC, mandatory=.false.)
      call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                           fname_ne, fname_sw, 'delz', var_bc=atm%neststruct%delz_BC, mandatory=.false.)
!                           fname_ne, fname_sw, 'delz', Atm%delz, Atm%neststruct%delz_BC, mandatory=.false.)
    endif
#ifdef USE_COND
       call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                            fname_ne, fname_sw,'q_con', var_bc=atm%neststruct%q_con_BC, mandatory=.false.)
#ifdef MOIST_CAPPA
       call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
            fname_ne, fname_sw, 'cappa', var_bc=atm%neststruct%cappa_BC, mandatory=.false.)
#endif
#endif
#endif
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'u', atm%u, atm%neststruct%u_BC, jstag=1)
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'v', atm%v, atm%neststruct%v_BC, istag=1)
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'uc', var_bc=atm%neststruct%uc_BC, istag=1)
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'vc', var_bc=atm%neststruct%vc_BC, jstag=1)
    call register_bcs_3d(atm, atm%neststruct%BCfile_ne, atm%neststruct%BCfile_sw, &
                         fname_ne, fname_sw, 'divg', var_bc=atm%neststruct%divg_BC, istag=1,jstag=1, mandatory=.false.)

    return
  end subroutine fv_io_register_restart_bcs


  subroutine fv_io_write_bcs(Atm, timestamp)
    type(fv_atmos_type), intent(inout) :: Atm
    character(len=*),    intent(in), optional :: timestamp

    call save_restart_border(atm%neststruct%BCfile_ne, timestamp)
    call save_restart_border(atm%neststruct%BCfile_sw, timestamp)

    return
  end subroutine fv_io_write_bcs


  subroutine fv_io_read_bcs(Atm)
    type(fv_atmos_type), intent(inout) :: Atm

    call restore_state_border(atm%neststruct%BCfile_ne)
    call restore_state_border(atm%neststruct%BCfile_sw)

    !These do not work yet
    !need to modify register_bcs_?d to get ids for registered variables, and then use query_initialized_id
    !Atm%neststruct%divg_BC%initialized = field_exist(fname_ne, 'divg_north_t1', Atm%domain)
    !Atm%neststruct%w_BC%initialized    = field_exist(fname_ne, 'w_north_t1', Atm%domain)
    !Atm%neststruct%delz_BC%initialized = field_exist(fname_ne, 'delz_north_t1', Atm%domain)
    !if (is_master()) print*, ' BCs: ', Atm%neststruct%divg_BC%initialized, Atm%neststruct%w_BC%initialized, Atm%neststruct%delz_BC%initialized

    return
  end subroutine fv_io_read_bcs

end module fv_io_mod