sfx_esm.f90

module sfx_esm
    !> @brief main Earth System Model surface flux module

    ! macro defs.
    ! --------------------------------------------------------------------------------
#define SFX_FORCE_DEPRECATED_CODE
    ! --------------------------------------------------------------------------------

    ! modules used
    ! --------------------------------------------------------------------------------
    use sfx_esm_param
    ! --------------------------------------------------------------------------------

    ! directives list
    ! --------------------------------------------------------------------------------
    implicit none
    private
    ! --------------------------------------------------------------------------------

    ! public interface
    ! --------------------------------------------------------------------------------
    public :: get_surface_fluxes
    public :: get_surface_fluxes_vec
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    type, public :: meteoDataType
        !> @brief meteorological input for surface flux calculation

        real :: h       !> constant flux layer height [m]
        real :: U       !> abs(wind speed) at 'h' [m/s]
        real :: dT      !> difference between potential temperature at 'h' and at surface [K]
        real :: Tsemi   !> semi-sum of potential temperature at 'h' and at surface [K]
        real :: dQ      !> difference between humidity at 'h' and at surface [g/g]
        real :: z0_m    !> surface aerodynamic roughness (should be < 0 for water bodies surface)
    end type

    type, public :: meteoDataVecType
        !> @brief meteorological input for surface flux calculation
        !> &details using arrays as input

        real, allocatable :: h(:)       !> constant flux layer height [m]
        real, allocatable :: U(:)       !> abs(wind speed) at 'h' [m/s]
        real, allocatable :: dT(:)      !> difference between potential temperature at 'h' and at surface [K]
        real, allocatable :: Tsemi(:)   !> semi-sum of potential temperature at 'h' and at surface [K]
        real, allocatable :: dQ(:)      !> difference between humidity at 'h' and at surface [g/g]
        real, allocatable :: z0_m(:)    !> surface aerodynamic roughness (should be < 0 for water bodies surface)
    end type
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    type, public :: sfxDataType
        !> @brief surface flux output data

        real :: zeta            !> = z/L [n/d]
        real :: Rib             !> bulk Richardson number [n/d]
        real :: Re              !> Reynolds number [n/d]
        real :: B               !> = log(z0_m / z0_h) [n/d]
        real :: z0_m            !> aerodynamic roughness [m]
        real :: z0_t            !> thermal roughness [m]
        real :: Rib_conv_lim    !> Ri-bulk convection critical value [n/d]
        real :: Cm              !> transfer coefficient for momentum [n/d]
        real :: Ct              !> transfer coefficient for heat [n/d]
        real :: Km              !> eddy viscosity coeff. at h [m^2/s]
        real :: Pr_t_inv        !> inverse turbulent Prandtl number at h [n/d]
    end type

    type, public :: sfxDataVecType
        !> @brief surface flux output data
        !> &details using arrays as output

        real, allocatable :: zeta(:)            !> = z/L [n/d]
        real, allocatable :: Rib(:)             !> bulk Richardson number [n/d]
        real, allocatable :: Re(:)              !> Reynolds number [n/d]
        real, allocatable :: B(:)               !> = log(z0_m / z0_h) [n/d]
        real, allocatable :: z0_m(:)            !> aerodynamic roughness [m]
        real, allocatable :: z0_t(:)            !> thermal roughness [m]
        real, allocatable :: Rib_conv_lim(:)    !> Ri-bulk convection critical value [n/d]
        real, allocatable :: Cm(:)              !> transfer coefficient for momentum [n/d]
        real, allocatable :: Ct(:)              !> transfer coefficient for heat [n/d]
        real, allocatable :: Km(:)              !> eddy viscosity coeff. at h [m^2/s]
        real, allocatable :: Pr_t_inv(:)        !> inverse turbulent Prandtl number at h [n/d]
    end type
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    type, public :: numericsType
        integer :: maxiters_convection = 10    !> maximum (actual) number of iterations in convection
        integer :: maxiters_charnock = 10      !> maximum (actual) number of iterations in charnock roughness
    end type
    ! --------------------------------------------------------------------------------

contains

    ! --------------------------------------------------------------------------------
    subroutine get_surface_fluxes_vec(sfx, meteo, numerics, n)
        !> @brief surface flux calculation for array data
        !> @details contains C/C++ & CUDA interface
        ! ----------------------------------------------------------------------------
        type (sfxDataVecType), intent(inout) :: sfx

        type (meteoDataVecType), intent(in) :: meteo
        type (numericsType), intent(in) :: numerics
        integer, intent(in) :: n
        ! ----------------------------------------------------------------------------

        ! --- local variables
        type (meteoDataType)  meteo_cell
        type (sfxDataType) sfx_cell
        integer i
        ! ----------------------------------------------------------------------------

        do i = 1, n
#ifdef USE_DEPRECATED_CODE
#else
            meteo_cell = meteoDataType(&
                    h = meteo%h(i), &
                    U = meteo%U(i), dT = meteo%dT(i), Tsemi = meteo%Tsemi(i), dQ = meteo%dQ(i), &
                    z0_m = meteo%z0_m(i))

            call get_surface_fluxes(sfx_cell, meteo_cell, numerics)

            call push_sfx_data(sfx, sfx_cell, i)
#endif
        end do

    end subroutine get_surface_fluxes_vec
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    subroutine push_sfx_data(sfx, sfx_cell, idx)
        !> @brief helper subroutine for copying data in sfxDataVecType
        ! ----------------------------------------------------------------------------
        type (sfxDataVecType), intent(inout) :: sfx
        type (sfxDataType), intent(in) :: sfx_cell
        integer, intent(in) :: idx
        ! ----------------------------------------------------------------------------

        sfx%zeta(idx) = sfx_cell%zeta
        sfx%Rib(idx) = sfx_cell%Rib
        sfx%Re(idx) = sfx_cell%Re
        sfx%B(idx) = sfx_cell%B
        sfx%z0_m(idx) = sfx_cell%z0_m
        sfx%z0_t(idx) = sfx_cell%z0_t
        sfx%Rib_conv_lim(idx) = sfx_cell%Rib_conv_lim
        sfx%Cm(idx) = sfx_cell%Cm
        sfx%Ct(idx) = sfx_cell%Ct
        sfx%Km(idx) = sfx_cell%Km
        sfx%Pr_t_inv(idx) = sfx_cell%Pr_t_inv

    end subroutine push_sfx_data
    ! --------------------------------------------------------------------------------

    ! --------------------------------------------------------------------------------
    subroutine get_surface_fluxes(sfx, meteo, numerics)
        !> @brief surface flux calculation for single cell
        !> @details contains C/C++ interface
        ! ----------------------------------------------------------------------------
        type (sfxDataType), intent(out) :: sfx

        type (meteoDataType), intent(in) :: meteo
        type (numericsType), intent(in) :: numerics
        ! ----------------------------------------------------------------------------

        ! --- meteo derived datatype name shadowing
        ! ----------------------------------------------------------------------------
        real :: h       !> constant flux layer height [m]
        real :: U       !> abs(wind speed) at 'h' [m/s]
        real :: dT      !> difference between potential temperature at 'h' and at surface [K]
        real :: Tsemi   !> semi-sum of potential temperature at 'h' and at surface [K]
        real :: dQ      !> difference between humidity at 'h' and at surface [g/g]
        real :: z0_m    !> surface aerodynamic roughness (should be < 0 for water bodies surface)
        ! ----------------------------------------------------------------------------

        ! --- local variables
        ! ----------------------------------------------------------------------------
        real z0_t               !> thermal roughness [m]
        real B                  !> = ln(z0_m / z0_t) [n/d]
        real h0_m, h0_t         !> = h / z0_m, h / z0_h [n/d]

        real u_dyn0             !> dynamic velocity in neutral conditions [m/s]
        real Re                 !> roughness Reynolds number = u_dyn0 * z0_m / nu [n/d]

        real zeta               !> = z/L [n/d]
        real Rib                !> bulk Richardson number

        real zeta_conv_lim      !> z/L critical value for matching free convection limit [n/d]
        real Rib_conv_lim       !> Ri-bulk critical value for matching free convection limit [n/d]

        real f_m_conv_lim       !> stability function (momentum) value in free convection limit [n/d]
        real f_h_conv_lim       !> stability function (heat) value in free convection limit [n/d]

        real psi_m, psi_h       !> universal functions (momentum) & (heat) [n/d]
        real phi_m, phi_h       !> stability functions (momentum) & (heat) [n/d]

        real Km                 !> eddy viscosity coeff. at h [m^2/s]
        real Pr_t_inv           !> invese Prandt number [n/d]

        real Cm, Ct             !> transfer coeff. for (momentum) & (heat) [n/d]

        integer surface_type    !> surface type = (ocean || land)

        real fval               !> just a shortcut for partial calculations
        ! ----------------------------------------------------------------------------

        ! --- shadowing names for clarity
        U = meteo%U
        Tsemi = meteo%Tsemi
        dT = meteo%dT
        dQ = meteo%dQ
        h = meteo%h
        z0_m = meteo%z0_m

        ! --- define surface type
        if (z0_m < 0.0) then
            surface_type = surface_ocean
        else
            surface_type = surface_land
        end if

        if (surface_type == surface_ocean) then
            ! --- define surface roughness [momentum] & dynamic velocity in neutral conditions
            call get_charnock_roughness(z0_m, u_dyn0, U, h, numerics%maxiters_charnock)
            ! --- define relative height
            h0_m = h / z0_m
        endif
        if (surface_type == surface_land) then
            ! --- define relative height
            h0_m = h / z0_m
            ! --- define dynamic velocity in neutral conditions
            u_dyn0 = U * kappa / log(h0_m)
        end if

        ! --- define B = log(z0_m / z0_h)
        Re = u_dyn0 * z0_m / nu_air
        if(Re <= Re_rough_min) then
            B = B1_rough * alog(B3_rough * Re) + B2_rough
        else
            ! *: B4 takes into account Re value at z' ~ O(10) z0
            B = B4_rough * (Re**B2_rough)
        end if

        !   --- apply max restriction based on surface type
        if (surface_type == surface_ocean) then
            B = min(B, B_max_ocean)
        endif
        if (surface_type == surface_land) then
            B = min(B, B_max_land)
        end if

        ! --- define roughness [thermal]
        z0_t = z0_m / exp(B)
        ! --- define relative height [thermal]
        h0_t = h / z0_t

        ! --- define Ri-bulk
        Rib = (g / Tsemi) * h * (dT + 0.61e0 * Tsemi * dQ) / U**2

        ! --- define free convection transition zeta = z/L value
        call get_convection_lim(zeta_conv_lim, Rib_conv_lim, f_m_conv_lim, f_h_conv_lim, &
                h0_m, h0_t, B)

        ! --- get the fluxes
        ! ----------------------------------------------------------------------------
        if (Rib > 0.0) then
            ! --- stable stratification block

            !   --- restrict bulk Ri value
            !   *: note that this value is written to output
            Rib = min(Rib, Rib_max)
            call get_psi_stable(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B)

            fval = beta_m * zeta
            phi_m = 1.0 + fval
            phi_h = 1.0/Pr_t_0_inv + fval

        else if (Rib < Rib_conv_lim) then
            ! --- strong instability block

            call get_psi_convection(psi_m, psi_h, zeta, Rib, &
                    zeta_conv_lim, f_m_conv_lim, f_h_conv_lim, h0_m, h0_t, B, numerics%maxiters_convection)

            fval = (zeta_conv_lim / zeta)**(1.0/3.0)
            phi_m = fval / f_m_conv_lim
            phi_h = fval / (Pr_t_0_inv * f_h_conv_lim)

        else if (Rib > -0.001) then
            ! --- nearly neutral [-0.001, 0] block

            call get_psi_neutral(psi_m, psi_h, zeta, h0_m, h0_t, B)

            phi_m = 1.0
            phi_h = 1.0 / Pr_t_0_inv
        else
            ! --- weak & semistrong instability block

            call get_psi_semi_convection(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B, numerics%maxiters_convection)

            phi_m = (1.0 - alpha_m * zeta)**(-0.25)
            phi_h = 1.0 / (Pr_t_0_inv * sqrt(1.0 - alpha_h_fix * zeta))
        end if

        ! --- define transfer coeff. (momentum) & (heat)
        Cm = kappa / psi_m
        Ct = kappa / psi_h

        ! --- define eddy viscosity & inverse Prandtl number
        Km = kappa * Cm * U * h / phi_m
        Pr_t_inv = phi_m / phi_h

        ! --- setting output
        sfx = sfxDataType(zeta = zeta, Rib = Rib, &
            Re = Re, B = B, z0_m = z0_m, z0_t = z0_t, &
            Rib_conv_lim = Rib_conv_lim, &
            Cm = Cm, Ct = Ct, Km = Km, Pr_t_inv = Pr_t_inv)

    end subroutine get_surface_fluxes
    ! --------------------------------------------------------------------------------

    ! convection universal functions shortcuts
    ! --------------------------------------------------------------------------------
    function f_m_conv(zeta)
        ! ----------------------------------------------------------------------------
        real :: f_m_conv
        real, intent(in) :: zeta
        ! ----------------------------------------------------------------------------

        f_m_conv = (1.0 - alpha_m * zeta)**0.25

    end function f_m_conv

    function f_h_conv(zeta)
        ! ----------------------------------------------------------------------------
        real :: f_h_conv
        real, intent(in) :: zeta
        ! ----------------------------------------------------------------------------

        f_h_conv = (1.0 - alpha_h * zeta)**0.5

    end function f_h_conv
    ! --------------------------------------------------------------------------------

    ! universal functions
    ! --------------------------------------------------------------------------------
    subroutine get_psi_neutral(psi_m, psi_h, zeta, h0_m, h0_t, B)
        !> @brief universal functions (momentum) & (heat): neutral case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h   !> universal functions
        real, intent(out) :: zeta           !> = z/L

        real, intent(in) :: h0_m, h0_t      !> = z/z0_m, z/z0_h
        real, intent(in) :: B               !> = log(z0_m / z0_h)
        ! ----------------------------------------------------------------------------

        zeta = 0.0
        psi_m = log(h0_m)
        psi_h = log(h0_t) / Pr_t_0_inv
        if (abs(B) < 1.0e-10) psi_h = psi_m / Pr_t_0_inv

    end subroutine

    subroutine get_psi_stable(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B)
        !> @brief universal functions (momentum) & (heat): stable case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h   !> universal functions [n/d]
        real, intent(out) :: zeta           !> = z/L [n/d]

        real, intent(in) :: Rib             !> bulk Richardson number [n/d]
        real, intent(in) :: h0_m, h0_t      !> = z/z0_m, z/z0_h [n/d]
        real, intent(in) :: B               !> = log(z0_m / z0_h) [n/d]
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: Rib_coeff
        real :: psi0_m, psi0_h
        real :: phi
        real :: c
        ! ----------------------------------------------------------------------------

        psi0_m = alog(h0_m)
        psi0_h = B / psi0_m

        Rib_coeff = beta_m * Rib
        c = (psi0_h + 1.0) / Pr_t_0_inv - 2.0 * Rib_coeff
        zeta = psi0_m * (sqrt(c**2 + 4.0 * Rib_coeff * (1.0 - Rib_coeff)) - c) / &
                (2.0 * beta_m * (1.0 - Rib_coeff))

        phi = beta_m * zeta

        psi_m = psi0_m + phi
        psi_h = (psi0_m + B) / Pr_t_0_inv + phi

    end subroutine

    subroutine get_psi_semi_convection(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B, maxiters)
        !> @brief universal functions (momentum) & (heat): semi-strong convection case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h   !> universal functions [n/d]
        real, intent(out) :: zeta           !> = z/L [n/d]

        real, intent(in) :: Rib             !> bulk Richardson number [n/d]
        real, intent(in) :: h0_m, h0_t      !> = z/z0_m, z/z0_h [n/d]
        real, intent(in) :: B               !> = log(z0_m / z0_h) [n/d]
        integer, intent(in) :: maxiters     !> maximum number of iterations

        ! --- local variables
        real :: zeta0_m, zeta0_h
        real :: f0_m, f0_h
        real :: f_m, f_h

        integer :: i
        ! ----------------------------------------------------------------------------

        psi_m = log(h0_m)
        psi_h = log(h0_t)
        if (abs(B) < 1.0e-10) psi_h = psi_m

        zeta = Rib * Pr_t_0_inv * psi_m**2 / psi_h

        do i = 1, maxiters + 1
            zeta0_m = zeta / h0_m
            zeta0_h = zeta / h0_t
            if (abs(B) < 1.0e-10) zeta0_h = zeta0_m

            f_m = (1.0 - alpha_m * zeta)**0.25e0
            f_h = sqrt(1.0 - alpha_h_fix * zeta)

            f0_m = (1.0 - alpha_m * zeta0_m)**0.25e0
            f0_h = sqrt(1.0 - alpha_h_fix * zeta0_h)

            f0_m = max(f0_m, 1.000001e0)
            f0_h = max(f0_h, 1.000001e0)

            psi_m = log((f_m - 1.0e0)*(f0_m + 1.0e0)/((f_m + 1.0e0)*(f0_m - 1.0e0))) + 2.0e0*(atan(f_m) - atan(f0_m))
            psi_h = log((f_h - 1.0e0)*(f0_h + 1.0e0)/((f_h + 1.0e0)*(f0_h - 1.0e0))) / Pr_t_0_inv

            ! *: don't update zeta = z/L at last iteration
            if (i == maxiters + 1) exit

            zeta = Rib * psi_m**2 / psi_h
        end do

    end subroutine

    subroutine get_psi_convection(psi_m, psi_h, zeta, Rib, &
            zeta_conv_lim, f_m_conv_lim, f_h_conv_lim, &
            h0_m, h0_t, B, maxiters)
        !> @brief universal functions (momentum) & (heat): fully convective case
        ! ----------------------------------------------------------------------------
        real, intent(out) :: psi_m, psi_h               !> universal functions [n/d]
        real, intent(out) :: zeta                       !> = z/L [n/d]

        real, intent(in) :: Rib                         !> bulk Richardson number [n/d]
        real, intent(in) :: h0_m, h0_t                  !> = z/z0_m, z/z0_h [n/d]
        real, intent(in) :: B                           !> = log(z0_m / z0_h) [n/d]
        integer, intent(in) :: maxiters                 !> maximum number of iterations

        real, intent(in) :: zeta_conv_lim               !> convective limit zeta
        real, intent(in) :: f_m_conv_lim, f_h_conv_lim  !> universal function shortcuts in limiting case
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: zeta0_m, zeta0_h
        real :: f0_m, f0_h
        real :: p_m, p_h
        real :: a_m, a_h
        real :: c_lim, f

        integer :: i
        ! ----------------------------------------------------------------------------

        p_m = 2.0 * atan(f_m_conv_lim) + log((f_m_conv_lim - 1.0) / (f_m_conv_lim + 1.0))
        p_h = log((f_h_conv_lim - 1.0) / (f_h_conv_lim + 1.0))

        zeta = zeta_conv_lim
        do i = 1, maxiters + 1
            zeta0_m = zeta / h0_m
            zeta0_h = zeta / h0_t
            if (abs(B) < 1.0e-10) zeta0_h = zeta0_m

            f0_m = (1.0 - alpha_m * zeta0_m)**0.25
            f0_h = sqrt(1.0 - alpha_h_fix * zeta0_h)

            a_m = -2.0*atan(f0_m) + log((f0_m + 1.0)/(f0_m - 1.0))
            a_h = log((f0_h + 1.0)/(f0_h - 1.0))

            c_lim = (zeta_conv_lim / zeta)**(1.0 / 3.0)
            f = 3.0 * (1.0 - c_lim)

            psi_m = f / f_m_conv_lim + p_m + a_m
            psi_h = (f / f_h_conv_lim + p_h + a_h) / Pr_t_0_inv

            ! *: don't update zeta = z/L at last iteration
            if (i == maxiters + 1) exit

            zeta = Rib * psi_m**2 / psi_h
        end do

    end subroutine
    ! --------------------------------------------------------------------------------

    ! convection limit definition
    ! --------------------------------------------------------------------------------
    subroutine get_convection_lim(zeta_lim, Rib_lim, f_m_lim, f_h_lim, &
            h0_m, h0_t, B)
        ! ----------------------------------------------------------------------------
        real, intent(out) :: zeta_lim           !> limiting value of z/L
        real, intent(out) :: Rib_lim            !> limiting value of Ri-bulk
        real, intent(out) :: f_m_lim, f_h_lim   !> limiting values of universal functions shortcuts

        real, intent(in) :: h0_m, h0_t          !> = z/z0_m, z/z0_h [n/d]
        real, intent(in) :: B                   !> = log(z0_m / z0_h) [n/d]
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: psi_m, psi_h
        real :: f_m, f_h
        real :: c
        ! ----------------------------------------------------------------------------

        ! --- define limiting value of zeta = z / L
        c = (Pr_t_inf_inv / Pr_t_0_inv)**4
        zeta_lim = (2.0 * alpha_h - c * alpha_m - &
                sqrt((c * alpha_m)**2 + 4.0 * c * alpha_h * (alpha_h - alpha_m))) / (2.0 * alpha_h**2)

        f_m_lim = f_m_conv(zeta_lim)
        f_h_lim = f_h_conv(zeta_lim)

        ! --- universal functions
        f_m = zeta_lim / h0_m
        f_h = zeta_lim / h0_t
        if (abs(B) < 1.0e-10) f_h = f_m

        f_m = (1.0 - alpha_m * f_m)**0.25
        f_h = sqrt(1.0 - alpha_h_fix * f_h)

        psi_m = 2.0 * (atan(f_m_lim) - atan(f_m)) + alog((f_m_lim - 1.0) * (f_m + 1.0)/((f_m_lim + 1.0) * (f_m - 1.0)))
        psi_h = alog((f_h_lim - 1.0) * (f_h + 1.0)/((f_h_lim + 1.0) * (f_h - 1.0))) / Pr_t_0_inv

        ! --- bulk Richardson number
        Rib_lim = zeta_lim * psi_h / (psi_m * psi_m)

    end subroutine
    ! --------------------------------------------------------------------------------

    ! charnock roughness definition
    ! --------------------------------------------------------------------------------
    subroutine get_charnock_roughness(z0_m, u_dyn0, U, h, maxiters)
        ! ----------------------------------------------------------------------------
        real, intent(out) :: z0_m           !> aerodynamic roughness [m]
        real, intent(out) :: u_dyn0         !> dynamic velocity in neutral conditions [m/s]

        real, intent(in) :: h               !> constant flux layer height [m]
        real, intent(in) :: U               !> abs(wind speed) [m/s]
        integer, intent(in) :: maxiters     !> maximum number of iterations
        ! ----------------------------------------------------------------------------

        ! --- local variables
        real :: Uc                  ! wind speed at h_charnock [m/s]

        real :: a, b, c, c_min
        real :: f

        integer :: i, j
        ! ----------------------------------------------------------------------------

        Uc = U
        a = 0.0
        b = 25.0
        c_min = log(h_charnock) / kappa

        do i = 1, maxiters
            f = c1_charnock - 2.0 * log(Uc)
            do j = 1, maxiters
                c = (f + 2.0 * log(b)) / kappa
                ! looks like the check should use U10 instead of U
                !   but note that a1 should be set = 0 in cycle beforehand
                if (U <= 8.0e0) a = log(1.0 + c2_charnock * ((b / Uc)**3)) / kappa
                c = max(c - a, c_min)
                b = c
            end do
            z0_m = h_charnock * exp(-c * kappa)
            z0_m = max(z0_m, 0.000015e0)
            Uc = U * log(h_charnock / z0_m) / log(h / z0_m)
        end do

        ! --- define dynamic velocity in neutral conditions
        u_dyn0 = Uc / c

    end subroutine
    ! --------------------------------------------------------------------------------

end module sfx_esm