tracer_example.F90

!> A sample tracer package that has striped initial conditions
module user_tracer_example

! This file is part of MOM6. See LICENSE.md for the license.

use mom_diag_mediator, only : diag_ctrl
use mom_error_handler, only : mom_error, fatal, warning
use mom_file_parser, only : get_param, log_param, log_version, param_file_type
use mom_forcing_type, only : forcing
use mom_grid, only : ocean_grid_type
use mom_hor_index, only : hor_index_type
use mom_io, only : file_exists, mom_read_data, slasher, vardesc, var_desc, query_vardesc
use mom_open_boundary, only : ocean_obc_type
use mom_restart, only : mom_restart_cs
use mom_sponge, only : set_up_sponge_field, sponge_cs
use mom_time_manager, only : time_type
use mom_tracer_registry, only : register_tracer, tracer_registry_type
use mom_unit_scaling, only : unit_scale_type
use mom_variables, only : surface
use mom_verticalgrid, only : verticalgrid_type

use coupler_types_mod, only : coupler_type_set_data, ind_csurf
use atmos_ocean_fluxes_mod, only : aof_set_coupler_flux

implicit none ; private

#include <MOM_memory.h>

public user_register_tracer_example, user_initialize_tracer, user_tracer_stock
public tracer_column_physics, user_tracer_surface_state, user_tracer_example_end

integer, parameter :: ntr = 1 !< The number of tracers in this module.

!> The control structure for the USER_tracer_example module
type, public :: user_tracer_example_cs ; private
  logical :: coupled_tracers = .false. !< These tracers are not offered to the coupler.
  character(len=200) :: tracer_ic_file !< The full path to the IC file, or " "
                                       !! to initialize internally.
  type(time_type), pointer :: time => null() !< A pointer to the ocean model's clock.
  type(tracer_registry_type), pointer :: tr_reg => null() !< A pointer to the tracer registry
  real, pointer :: tr(:,:,:,:) => null()  !< The array of tracers used in this subroutine, in g m-3?
  real :: land_val(ntr) = -1.0 !< The value of tr that is used where land is masked out.
  logical :: use_sponge    !< If true, sponges may be applied somewhere in the domain.

  integer, dimension(NTR) :: ind_tr !< Indices returned by aof_set_coupler_flux if it is used and the
                                    !! surface tracer concentrations are to be provided to the coupler.

  type(diag_ctrl), pointer :: diag => null() !< A structure that is used to regulate the timing of diagnostic output.

  type(vardesc) :: tr_desc(ntr) !< Descriptions of each of the tracers.
end type user_tracer_example_cs

contains

!> This subroutine is used to register tracer fields and subroutines to be used with MOM.
function user_register_tracer_example(HI, GV, param_file, CS, tr_Reg, restart_CS)
  type(hor_index_type),    intent(in)   :: hi   !< A horizontal index type structure
  type(verticalgrid_type), intent(in)   :: gv   !< The ocean's vertical grid structure
  type(param_file_type),   intent(in)   :: param_file !< A structure to parse for run-time parameters
  type(user_tracer_example_cs), pointer :: cs   !< A pointer that is set to point to the control
                                                !! structure for this module
  type(tracer_registry_type), pointer   :: tr_reg !< A pointer that is set to point to the control
                                                  !! structure for the tracer advection and
                                                  !! diffusion module
  type(mom_restart_cs),       pointer   :: restart_cs !< A pointer to the restart control structure

! Local variables
  character(len=80)  :: name, longname
! This include declares and sets the variable "version".
#include "version_variable.h"
  character(len=40)  :: mdl = "tracer_example" ! This module's name.
  character(len=200) :: inputdir
  character(len=48) :: flux_units ! The units for tracer fluxes, usually
                            ! kg(tracer) kg(water)-1 m3 s-1 or kg(tracer) s-1.
  real, pointer :: tr_ptr(:,:,:) => null()
  logical :: user_register_tracer_example
  integer :: isd, ied, jsd, jed, nz, m
  isd = hi%isd ; ied = hi%ied ; jsd = hi%jsd ; jed = hi%jed ; nz = gv%ke

  if (associated(cs)) then
    call mom_error(warning, "USER_register_tracer_example called with an "// &
                            "associated control structure.")
    return
  endif
  allocate(cs)

  ! Read all relevant parameters and write them to the model log.
  call log_version(param_file, mdl, version, "")
  call get_param(param_file, mdl, "TRACER_EXAMPLE_IC_FILE", cs%tracer_IC_file, &
                 "The name of a file from which to read the initial "//&
                 "conditions for the DOME tracers, or blank to initialize "//&
                 "them internally.", default=" ")
  if (len_trim(cs%tracer_IC_file) >= 1) then
    call get_param(param_file, mdl, "INPUTDIR", inputdir, default=".")
    cs%tracer_IC_file = trim(slasher(inputdir))//trim(cs%tracer_IC_file)
    call log_param(param_file, mdl, "INPUTDIR/TRACER_EXAMPLE_IC_FILE", &
                   cs%tracer_IC_file)
  endif
  call get_param(param_file, mdl, "SPONGE", cs%use_sponge, &
                 "If true, sponges may be applied anywhere in the domain. "//&
                 "The exact location and properties of those sponges are "//&
                 "specified from MOM_initialization.F90.", default=.false.)

  allocate(cs%tr(isd:ied,jsd:jed,nz,ntr)) ; cs%tr(:,:,:,:) = 0.0

  do m=1,ntr
    if (m < 10) then ; write(name,'("tr",I1.1)') m
    else ; write(name,'("tr",I2.2)') m ; endif
    write(longname,'("Concentration of Tracer ",I2.2)') m
    cs%tr_desc(m) = var_desc(name, units="kg kg-1", longname=longname, caller=mdl)

    ! This needs to be changed if the units of tracer are changed above.
    if (gv%Boussinesq) then ; flux_units = "kg kg-1 m3 s-1"
    else ; flux_units = "kg s-1" ; endif

    ! This is needed to force the compiler not to do a copy in the registration
    ! calls.  Curses on the designers and implementers of Fortran90.
    tr_ptr => cs%tr(:,:,:,m)
    ! Register the tracer for horizontal advection, diffusion, and restarts.
    call register_tracer(tr_ptr, tr_reg, param_file, hi, gv, &
                         name=name, longname=longname, units="kg kg-1", &
                         registry_diags=.true., flux_units=flux_units, &
                         restart_cs=restart_cs)

    !   Set coupled_tracers to be true (hard-coded above) to provide the surface
    ! values to the coupler (if any).  This is meta-code and its arguments will
    ! currently (deliberately) give fatal errors if it is used.
    if (cs%coupled_tracers) &
      cs%ind_tr(m) = aof_set_coupler_flux(trim(name)//'_flux', &
          flux_type=' ', implementation=' ', caller="USER_register_tracer_example")
  enddo

  cs%tr_Reg => tr_reg
  user_register_tracer_example = .true.
end function user_register_tracer_example

!> This subroutine initializes the NTR tracer fields in tr(:,:,:,:)
!! and it sets up the tracer output.
subroutine user_initialize_tracer(restart, day, G, GV, h, diag, OBC, CS, &
                                  sponge_CSp)
  logical,                            intent(in) :: restart !< .true. if the fields have already
                                                         !! been read from a restart file.
  type(time_type),            target, intent(in) :: day  !< Time of the start of the run.
  type(ocean_grid_type),              intent(in) :: g    !< The ocean's grid structure
  type(verticalgrid_type),            intent(in) :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                                      intent(in) :: h    !< Layer thicknesses [H ~> m or kg m-2]
  type(diag_ctrl),            target, intent(in) :: diag !< A structure that is used to regulate
                                                         !! diagnostic output.
  type(ocean_obc_type),               pointer    :: obc  !< This open boundary condition type specifies
                                                         !! whether, where, and what open boundary
                                                         !! conditions are used.
  type(user_tracer_example_cs),       pointer    :: cs   !< The control structure returned by a previous
                                                         !! call to USER_register_tracer_example.
  type(sponge_cs),                    pointer    :: sponge_csp    !< A pointer to the control structure
                                                                  !! for the sponges, if they are in use.

! Local variables
  real, allocatable :: temp(:,:,:)
  character(len=32) :: name     ! A variable's name in a NetCDF file.
  character(len=72) :: longname ! The long name of that variable.
  character(len=48) :: units    ! The dimensions of the variable.
  character(len=48) :: flux_units ! The units for tracer fluxes, usually
                            ! kg(tracer) kg(water)-1 m3 s-1 or kg(tracer) s-1.
  real, pointer :: tr_ptr(:,:,:) => null()
  real :: pi     ! 3.1415926... calculated as 4*atan(1)
  real :: tr_y   ! Initial zonally uniform tracer concentrations.
  real :: dist2  ! The distance squared from a line [m2].
  integer :: i, j, k, is, ie, js, je, isd, ied, jsd, jed, nz, m
  integer :: isdb, iedb, jsdb, jedb, lntr

  if (.not.associated(cs)) return
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke
  isd = g%isd ; ied = g%ied ; jsd = g%jsd ; jed = g%jed
  isdb = g%IsdB ; iedb = g%IedB ; jsdb = g%JsdB ; jedb = g%JedB

  lntr = ntr ! Avoids compile-time warning when NTR<2
  cs%Time => day
  cs%diag => diag

  if (.not.restart) then
    if (len_trim(cs%tracer_IC_file) >= 1) then
!  Read the tracer concentrations from a netcdf file.
      if (.not.file_exists(cs%tracer_IC_file, g%Domain)) &
        call mom_error(fatal, "USER_initialize_tracer: Unable to open "// &
                        cs%tracer_IC_file)
      do m=1,ntr
        call query_vardesc(cs%tr_desc(m), name, caller="USER_initialize_tracer")
        call mom_read_data(cs%tracer_IC_file, trim(name), cs%tr(:,:,:,m), g%Domain)
      enddo
    else
      do m=1,ntr
        do k=1,nz ; do j=js,je ; do i=is,ie
          cs%tr(i,j,k,m) = 1.0e-20 ! This could just as well be 0.
        enddo ; enddo ; enddo
      enddo

!    This sets a stripe of tracer across the basin.
      pi = 4.0*atan(1.0)
      do j=js,je
        dist2 = (g%Rad_Earth * pi / 180.0)**2 * &
               (g%geoLatT(i,j) - 40.0) * (g%geoLatT(i,j) - 40.0)
        tr_y = 0.5*exp(-dist2/(1.0e5*1.0e5))

        do k=1,nz ; do i=is,ie
!      This adds the stripes of tracer to every layer.
          cs%tr(i,j,k,1) = cs%tr(i,j,k,1) + tr_y
        enddo ; enddo
      enddo
    endif
  endif ! restart

  if ( cs%use_sponge ) then
!   If sponges are used, this example damps tracers in sponges in the
! northern half of the domain to 1 and tracers in the southern half
! to 0.  For any tracers that are not damped in the sponge, the call
! to set_up_sponge_field can simply be omitted.
    if (.not.associated(sponge_csp)) &
      call mom_error(fatal, "USER_initialize_tracer: "// &
        "The pointer to sponge_CSp must be associated if SPONGE is defined.")

    allocate(temp(g%isd:g%ied,g%jsd:g%jed,nz))
    do k=1,nz ; do j=js,je ; do i=is,ie
      if (g%geoLatT(i,j) > 700.0 .and. (k > nz/2)) then
        temp(i,j,k) = 1.0
      else
        temp(i,j,k) = 0.0
      endif
    enddo ; enddo ; enddo

!   do m=1,NTR
    do m=1,1
      ! This is needed to force the compiler not to do a copy in the sponge
      ! calls.  Curses on the designers and implementers of Fortran90.
      tr_ptr => cs%tr(:,:,:,m)
      call set_up_sponge_field(temp, tr_ptr, g, nz, sponge_csp)
    enddo
    deallocate(temp)
  endif

  if (associated(obc)) then
    call query_vardesc(cs%tr_desc(1), name, caller="USER_initialize_tracer")
    if (obc%specified_v_BCs_exist_globally) then
      ! Steal from updated DOME in the fullness of time.
    else
      ! Steal from updated DOME in the fullness of time.
    endif
    ! All tracers but the first have 0 concentration in their inflows. As this
    ! is the default value, the following calls are unnecessary.
    do m=2,lntr
      call query_vardesc(cs%tr_desc(m), name, caller="USER_initialize_tracer")
      ! Steal from updated DOME in the fullness of time.
    enddo
  endif

end subroutine user_initialize_tracer

!> This subroutine applies diapycnal diffusion and any other column
!! tracer physics or chemistry to the tracers from this file.
!! This is a simple example of a set of advected passive tracers.
!! The arguments to this subroutine are redundant in that
!!     h_new(k) = h_old(k) + ea(k) - eb(k-1) + eb(k) - ea(k+1)
subroutine tracer_column_physics(h_old, h_new,  ea,  eb, fluxes, dt, G, GV, US, CS)
  type(ocean_grid_type),   intent(in) :: g    !< The ocean's grid structure
  type(verticalgrid_type), intent(in) :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                           intent(in) :: h_old !< Layer thickness before entrainment [H ~> m or kg m-2].
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                           intent(in) :: h_new !< Layer thickness after entrainment [H ~> m or kg m-2].
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                           intent(in) :: ea   !< an array to which the amount of fluid entrained
                                              !! from the layer above during this call will be
                                              !! added [H ~> m or kg m-2].
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                           intent(in) :: eb   !< an array to which the amount of fluid entrained
                                              !! from the layer below during this call will be
                                              !! added [H ~> m or kg m-2].
  type(forcing),           intent(in) :: fluxes !< A structure containing pointers to thermodynamic
                                              !! and tracer forcing fields.  Unused fields have NULL ptrs.
  real,                    intent(in) :: dt   !< The amount of time covered by this call [T ~> s]
  type(unit_scale_type),   intent(in) :: us   !< A dimensional unit scaling type
  type(user_tracer_example_cs), pointer :: cs !< The control structure returned by a previous
                                              !! call to USER_register_tracer_example.

! Local variables
  real :: hold0(szi_(g))       ! The original topmost layer thickness,
                               ! with surface mass fluxes added back, m.
  real :: b1(szi_(g))          ! b1 and c1 are variables used by the
  real :: c1(szi_(g),szk_(g))  ! tridiagonal solver.
  real :: d1(szi_(g))          ! d1=1-c1 is used by the tridiagonal solver.
  real :: h_neglect            ! A thickness that is so small it is usually lost
                               ! in roundoff and can be neglected [H ~> m or kg m-2].
  real :: b_denom_1 ! The first term in the denominator of b1 [H ~> m or kg m-2].
  integer :: i, j, k, is, ie, js, je, nz, m

! The following array (trdc) determines the behavior of the tracer
! diapycnal advection.  The first element is 1 if tracers are
! passively advected.  The second and third are the concentrations
! to which downwelling and upwelling water are set, respectively.
! For most (normal) tracers, the appropriate vales are {1,0,0}.

  real :: trdc(3)
!   Uncomment the following line to dye both upwelling and downwelling.
! data trdc / 0.0,1.0,1.0 /
!   Uncomment the following line to dye downwelling.
! data trdc / 0.0,1.0,0.0 /
!   Uncomment the following line to dye upwelling.
! data trdc / 0.0,0.0,1.0 /
!   Uncomment the following line for tracer concentrations to be set
! to zero in any diapycnal motions.
! data trdc / 0.0,0.0,0.0 /
!   Uncomment the following line for most "physical" tracers, which
! are advected diapycnally in the usual manner.
  data trdc / 1.0,0.0,0.0 /
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke

  if (.not.associated(cs)) return
  h_neglect = gv%H_subroundoff

  do j=js,je
    do i=is,ie
!   The following line is appropriate for quantities like salinity
! that are left behind by evaporation, and any surface fluxes would
! be explicitly included in the flux structure.
      hold0(i) = h_old(i,j,1)
!   The following line is appropriate for quantities like temperature
! that can be assumed to have the same concentration in evaporation
! as they had in the water.  The explicit surface fluxes here would
! reflect differences in concentration from the ambient water, not
! the absolute fluxes.
  !   hold0(i) = h_old(i,j,1) + ea(i,j,1)
      b_denom_1 = h_old(i,j,1) + ea(i,j,1) + h_neglect
      b1(i) = 1.0 / (b_denom_1 + eb(i,j,1))
!       d1(i) = b_denom_1 * b1(i)
      d1(i) = trdc(1) * (b_denom_1 * b1(i)) + (1.0 - trdc(1))
      do m=1,ntr
        cs%tr(i,j,1,m) = b1(i)*(hold0(i)*cs%tr(i,j,1,m) + trdc(3)*eb(i,j,1))
 !      Add any surface tracer fluxes to the preceding line.
      enddo
    enddo
    do k=2,nz ; do i=is,ie
      c1(i,k) = trdc(1) * eb(i,j,k-1) * b1(i)
      b_denom_1 = h_old(i,j,k) + d1(i)*ea(i,j,k) + h_neglect
      b1(i) = 1.0 / (b_denom_1 + eb(i,j,k))
      d1(i) = trdc(1) * (b_denom_1 * b1(i)) + (1.0 - trdc(1))
      do m=1,ntr
        cs%tr(i,j,k,m) = b1(i) * (h_old(i,j,k)*cs%tr(i,j,k,m) + &
                 ea(i,j,k)*(trdc(1)*cs%tr(i,j,k-1,m)+trdc(2)) + &
                 eb(i,j,k)*trdc(3))
      enddo
    enddo ; enddo
    do m=1,ntr ; do k=nz-1,1,-1 ; do i=is,ie
      cs%tr(i,j,k,m) = cs%tr(i,j,k,m) + c1(i,k+1)*cs%tr(i,j,k+1,m)
    enddo ; enddo ; enddo
  enddo

end subroutine tracer_column_physics

!> This function calculates the mass-weighted integral of all tracer stocks,
!! returning the number of stocks it has calculated.  If the stock_index
!! is present, only the stock corresponding to that coded index is returned.
function user_tracer_stock(h, stocks, G, GV, CS, names, units, stock_index)
  type(ocean_grid_type),              intent(in)    :: g    !< The ocean's grid structure
  type(verticalgrid_type),            intent(in)    :: gv   !< The ocean's vertical grid structure
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                                      intent(in)    :: h    !< Layer thicknesses [H ~> m or kg m-2]
  real, dimension(:),                 intent(out)   :: stocks !< the mass-weighted integrated amount of each
                                                              !! tracer, in kg times concentration units [kg conc].
  type(user_tracer_example_cs),       pointer       :: cs     !< The control structure returned by a
                                                              !! previous call to register_USER_tracer.
  character(len=*), dimension(:),     intent(out)   :: names  !< The names of the stocks calculated.
  character(len=*), dimension(:),     intent(out)   :: units  !< The units of the stocks calculated.
  integer, optional,                  intent(in)    :: stock_index !< The coded index of a specific stock
                                                              !! being sought.
  integer                                           :: user_tracer_stock !< Return value: the number of
                                                              !! stocks calculated here.

! Local variables
  integer :: i, j, k, is, ie, js, je, nz, m
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec ; nz = gv%ke

  user_tracer_stock = 0
  if (.not.associated(cs)) return

  if (present(stock_index)) then ; if (stock_index > 0) then
    ! Check whether this stock is available from this routine.

    ! No stocks from this routine are being checked yet.  Return 0.
    return
  endif ; endif

  do m=1,ntr
    call query_vardesc(cs%tr_desc(m), name=names(m), units=units(m), caller="USER_tracer_stock")
    units(m) = trim(units(m))//" kg"
    stocks(m) = 0.0
    do k=1,nz ; do j=js,je ; do i=is,ie
      stocks(m) = stocks(m) + cs%tr(i,j,k,m) * &
                             (g%mask2dT(i,j) * g%US%L_to_m**2*g%areaT(i,j) * h(i,j,k))
    enddo ; enddo ; enddo
    stocks(m) = gv%H_to_kg_m2 * stocks(m)
  enddo
  user_tracer_stock = ntr

end function user_tracer_stock

!> This subroutine extracts the surface fields from this tracer package that
!! are to be shared with the atmosphere in coupled configurations.
subroutine user_tracer_surface_state(sfc_state, h, G, CS)
  type(ocean_grid_type),        intent(in)    :: g     !< The ocean's grid structure
  type(surface),                intent(inout) :: sfc_state !< A structure containing fields that
                                                       !! describe the surface state of the ocean.
  real, dimension(SZI_(G),SZJ_(G),SZK_(G)), &
                                intent(in)    :: h  !< Layer thicknesses [H ~> m or kg m-2]
  type(user_tracer_example_cs), pointer       :: cs !< The control structure returned by a previous
                                                    !! call to register_USER_tracer.

  ! This particular tracer package does not report anything back to the coupler.
  ! The code that is here is just a rough guide for packages that would.

  integer :: m, is, ie, js, je, isd, ied, jsd, jed
  is = g%isc ; ie = g%iec ; js = g%jsc ; je = g%jec
  isd = g%isd ; ied = g%ied ; jsd = g%jsd ; jed = g%jed

  if (.not.associated(cs)) return

  if (cs%coupled_tracers) then
    do m=1,ntr
      !   This call loads the surface values into the appropriate array in the
      ! coupler-type structure.
      call coupler_type_set_data(cs%tr(:,:,1,m), cs%ind_tr(m), ind_csurf, &
                   sfc_state%tr_fields, idim=(/isd, is, ie, ied/), &
                   jdim=(/jsd, js, je, jed/) )
    enddo
  endif

end subroutine user_tracer_surface_state

!> Clean up allocated memory at the end.
subroutine user_tracer_example_end(CS)
  type(user_tracer_example_cs), pointer :: cs !< The control structure returned by a previous
                                              !! call to register_USER_tracer.
  integer :: m

  if (associated(cs)) then
    if (associated(cs%tr)) deallocate(cs%tr)
    deallocate(cs)
  endif
end subroutine user_tracer_example_end

!> \namespace user_tracer_example
!!
!!  Original by Robert Hallberg, 2002
!!
!!    This file contains an example of the code that is needed to set
!!  up and use a set (in this case one) of dynamically passive tracers.
!!
!!    A single subroutine is called from within each file to register
!!  each of the tracers for reinitialization and advection and to
!!  register the subroutine that initializes the tracers and set up
!!  their output and the subroutine that does any tracer physics or
!!  chemistry along with diapycnal mixing (included here because some
!!  tracers may float or swim vertically or dye diapycnal processes).
end module user_tracer_example