set_analytic_initial_condition Subroutine

    subroutine set_analytic_initial_condition()
        ! Module(s) from CAM-SIMA.
        use cam_logfile, only: debugout_debug
        ! Module(s) from CESM Share.
        use shr_kind_mod, only: kind_r8 => shr_kind_r8

        character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition'
        integer, allocatable :: global_grid_index(:)
        real(kind_r8), allocatable :: buffer_2d_real(:, :), buffer_3d_real(:, :, :)
        real(kind_r8), allocatable :: lat_rad(:), lon_rad(:)
        real(kind_r8), allocatable :: z_int(:, :)   ! Geometric height (m) at layer interfaces.
                                                    ! Dimension and vertical index orders follow CAM-SIMA convention.
        real(kind_dyn_mpas), pointer :: zgrid(:, :) ! Geometric height (m) at layer interfaces.
                                                    ! Dimension and vertical index orders follow MPAS convention.

        call dyn_debug_print(debugout_debug, subname // ' entered')

        call init_shared_variables()

        call set_mpas_state_u()
        call set_mpas_state_w()
        call set_mpas_state_scalars()
        call set_mpas_state_rho_theta()
        call set_mpas_state_rho_base_theta_base()

        call final_shared_variables()

        call dyn_debug_print(debugout_debug, subname // ' completed')
    contains
        !> Initialize variables that are shared and repeatedly used by the `set_mpas_state_*` internal subroutines.
        !> (KCW, 2024-05-13)
        subroutine init_shared_variables()
            ! Module(s) from CAM-SIMA.
            use cam_abortutils, only: check_allocate, endrun
            use cam_grid_support, only: cam_grid_get_latvals, cam_grid_get_lonvals, cam_grid_id
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve
            use dyn_procedure, only: reverse
            use dynconst, only: deg_to_rad
            use vert_coord, only: pverp

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::init_shared_variables'
            integer :: i
            integer :: ierr
            integer, pointer :: indextocellid(:)
            real(kind_r8), pointer :: lat_deg(:), lon_deg(:)

            call dyn_debug_print(debugout_verbose, 'Preparing to set analytic initial condition')

            nullify(zgrid)
            nullify(indextocellid)
            nullify(lat_deg, lon_deg)

            allocate(global_grid_index(ncells_solve), stat=ierr)
            call check_allocate(ierr, subname, 'global_grid_index(ncells_solve)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(indextocellid, 'mesh', 'indexToCellID')

            global_grid_index(:) = indextocellid(1:ncells_solve)

            nullify(indextocellid)

            allocate(lat_rad(ncells_solve), stat=ierr)
            call check_allocate(ierr, subname, 'lat_rad(ncells_solve)', 'dyn_comp', __LINE__)

            allocate(lon_rad(ncells_solve), stat=ierr)
            call check_allocate(ierr, subname, 'lon_rad(ncells_solve)', 'dyn_comp', __LINE__)

            ! "mpas_cell" is a registered grid name that is defined in `dyn_grid`.
            lat_deg => cam_grid_get_latvals(cam_grid_id('mpas_cell'))
            lon_deg => cam_grid_get_lonvals(cam_grid_id('mpas_cell'))

            if (.not. associated(lat_deg)) then
                call endrun('Failed to find variable "lat_deg"', subname, __LINE__)
            end if

            if (.not. associated(lon_deg)) then
                call endrun('Failed to find variable "lon_deg"', subname, __LINE__)
            end if

            lat_rad(:) = lat_deg(:) * deg_to_rad
            lon_rad(:) = lon_deg(:) * deg_to_rad

            nullify(lat_deg, lon_deg)

            allocate(z_int(ncells_solve, pverp), stat=ierr)
            call check_allocate(ierr, subname, 'z_int(ncells_solve, pverp)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(zgrid, 'mesh', 'zgrid')

            ! Vertical index order is reversed between CAM-SIMA and MPAS.
            do i = 1, ncells_solve
                z_int(i, :) = reverse(real(zgrid(:, i), kind_r8))
            end do
        end subroutine init_shared_variables

        !> Finalize variables that are shared and repeatedly used by the `set_mpas_state_*` internal subroutines.
        !> (KCW, 2024-05-13)
        subroutine final_shared_variables()
            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::final_shared_variables'

            deallocate(global_grid_index)
            deallocate(lat_rad, lon_rad)
            deallocate(z_int)

            nullify(zgrid)
        end subroutine final_shared_variables

        !> Set MPAS state `u` (i.e., horizontal velocity at edge interfaces).
        !> (KCW, 2024-05-13)
        subroutine set_mpas_state_u()
            ! Module(s) from CAM-SIMA.
            use cam_abortutils, only: check_allocate
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve
            use dyn_procedure, only: reverse
            use dyn_tests_utils, only: vc_height
            use inic_analytic, only: dyn_set_inic_col
            use vert_coord, only: pver

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::set_mpas_state_u'
            integer :: i
            integer :: ierr
            real(kind_dyn_mpas), pointer :: ucellzonal(:, :), ucellmeridional(:, :)

            call dyn_debug_print(debugout_verbose, 'Setting MPAS state "u"')

            nullify(ucellzonal, ucellmeridional)

            allocate(buffer_2d_real(ncells_solve, pver), stat=ierr)
            call check_allocate(ierr, subname, 'buffer_2d_real(ncells_solve, pver)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(ucellzonal, 'diag', 'uReconstructZonal')
            call mpas_dynamical_core % get_variable_pointer(ucellmeridional, 'diag', 'uReconstructMeridional')

            buffer_2d_real(:, :) = 0.0_kind_r8

            call dyn_set_inic_col(vc_height, lat_rad, lon_rad, global_grid_index, zint=z_int, u=buffer_2d_real)

            ! Vertical index order is reversed between CAM-SIMA and MPAS.
            do i = 1, ncells_solve
                ucellzonal(:, i) = real(reverse(buffer_2d_real(i, :)), kind_dyn_mpas)
            end do

            buffer_2d_real(:, :) = 0.0_kind_r8

            call dyn_set_inic_col(vc_height, lat_rad, lon_rad, global_grid_index, zint=z_int, v=buffer_2d_real)

            ! Vertical index order is reversed between CAM-SIMA and MPAS.
            do i = 1, ncells_solve
                ucellmeridional(:, i) = real(reverse(buffer_2d_real(i, :)), kind_dyn_mpas)
            end do

            deallocate(buffer_2d_real)

            nullify(ucellzonal, ucellmeridional)

            call mpas_dynamical_core % compute_edge_wind(wind_tendency=.false.)
        end subroutine set_mpas_state_u

        !> Set MPAS state `w` (i.e., vertical velocity at cell interfaces).
        !> (KCW, 2024-05-13)
        subroutine set_mpas_state_w()
            ! Module(s) from CAM-SIMA.
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::set_mpas_state_w'
            real(kind_dyn_mpas), pointer :: w(:, :)

            call dyn_debug_print(debugout_verbose, 'Setting MPAS state "w"')

            nullify(w)

            call mpas_dynamical_core % get_variable_pointer(w, 'state', 'w', time_level=1)

            w(:, 1:ncells_solve) = 0.0_kind_dyn_mpas

            nullify(w)

            ! Because we are injecting data directly into MPAS memory, halo layers need to be updated manually.
            call mpas_dynamical_core % exchange_halo('w')
        end subroutine set_mpas_state_w

        !> Set MPAS state `scalars` (i.e., constituents).
        !> (KCW, 2024-05-17)
        subroutine set_mpas_state_scalars()
            ! Module(s) from CAM-SIMA.
            use cam_abortutils, only: check_allocate
            use cam_constituents, only: num_advected
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve
            use dyn_procedure, only: reverse
            use dyn_tests_utils, only: vc_height
            use inic_analytic, only: dyn_set_inic_col
            use vert_coord, only: pver

            ! CCPP standard name of `qv`, which denotes water vapor mixing ratio.
            character(*), parameter :: constituent_qv_standard_name = &
                'water_vapor_mixing_ratio_wrt_dry_air'

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::set_mpas_state_scalars'
            integer :: i, j
            integer :: ierr
            integer, allocatable :: constituent_index(:)
            integer, pointer :: index_qv
            real(kind_dyn_mpas), pointer :: scalars(:, :, :)

            call dyn_debug_print(debugout_verbose, 'Setting MPAS state "scalars"')

            nullify(index_qv)
            nullify(scalars)

            allocate(buffer_3d_real(ncells_solve, pver, num_advected), stat=ierr)
            call check_allocate(ierr, subname, 'buffer_3d_real(ncells_solve, pver, num_advected)', 'dyn_comp', __LINE__)

            allocate(constituent_index(num_advected), stat=ierr)
            call check_allocate(ierr, subname, 'constituent_index(num_advected)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(index_qv, 'dim', 'index_qv')
            call mpas_dynamical_core % get_variable_pointer(scalars, 'state', 'scalars', time_level=1)

            buffer_3d_real(:, :, :) = 0.0_kind_r8
            constituent_index(:) = [(i, i = 1, num_advected)]

            call dyn_set_inic_col(vc_height, lat_rad, lon_rad, global_grid_index, zint=z_int, q=buffer_3d_real, &
                m_cnst=constituent_index)

            do i = 1, ncells_solve
                ! `j` is indexing into `scalars`, so it is regarded as MPAS scalar index.
                do j = 1, num_advected
                    ! Vertical index order is reversed between CAM-SIMA and MPAS.
                    scalars(j, :, i) = &
                        real(reverse(buffer_3d_real(i, :, mpas_dynamical_core % map_constituent_index(j))), kind_dyn_mpas)
                end do
            end do

            if (mpas_dynamical_core % get_constituent_name(mpas_dynamical_core % map_constituent_index(index_qv)) == &
                constituent_qv_standard_name) then
                ! The definition of `qv` matches exactly what MPAS wants. No conversion is needed.
                call dyn_debug_print(debugout_verbose, 'No conversion is needed for water vapor mixing ratio')
            else
                ! The definition of `qv` actually represents specific humidity. Conversion is needed.
                call dyn_debug_print(debugout_verbose, 'Conversion is needed and applied for water vapor mixing ratio')

                ! Convert specific humidity to water vapor mixing ratio.
                scalars(index_qv, :, 1:ncells_solve) = &
                    scalars(index_qv, :, 1:ncells_solve) / (1.0_kind_dyn_mpas - scalars(index_qv, :, 1:ncells_solve))
            end if

            deallocate(buffer_3d_real)
            deallocate(constituent_index)

            nullify(index_qv)
            nullify(scalars)

            ! Because we are injecting data directly into MPAS memory, halo layers need to be updated manually.
            call mpas_dynamical_core % exchange_halo('scalars')
        end subroutine set_mpas_state_scalars

        !> Set MPAS state `rho` (i.e., dry air density) and `theta` (i.e., potential temperature).
        !> (KCW, 2024-05-19)
        subroutine set_mpas_state_rho_theta()
            ! Module(s) from CAM-SIMA.
            use cam_abortutils, only: check_allocate
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve
            use dyn_procedure, only: p_by_hypsometric_equation, rho_by_equation_of_state, theta_by_poisson_equation, &
                                     tm_of_t_qv, tv_of_tm_qv, reverse
            use dyn_tests_utils, only: vc_height
            use dynconst, only: constant_cpd => cpair, constant_g => gravit, constant_p0 => pref, &
                                constant_rd => rair, constant_rv => rh2o
            use inic_analytic, only: dyn_set_inic_col
            use vert_coord, only: pver

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::set_mpas_state_rho_theta'
            integer :: i, k
            integer :: ierr
            integer, pointer :: index_qv
            real(kind_r8), allocatable :: p_mid_col(:)  ! Pressure (Pa) at layer midpoints of each column. This is full pressure,
                                                        ! which also accounts for water vapor.
            real(kind_r8), allocatable :: p_sfc(:)      ! Pressure (Pa) at surface. This is full pressure,
                                                        ! which also accounts for water vapor.
            real(kind_r8), allocatable :: qv_mid_col(:) ! Water vapor mixing ratio (kg/kg) at layer midpoints of each column.
            real(kind_r8), allocatable :: t_mid(:, :)   ! Temperature (K) at layer midpoints.
            real(kind_r8), allocatable :: tm_mid_col(:) ! Modified "moist" temperature (K) at layer midpoints of each column.
                                                        ! Be advised that it is not virtual temperature.
                                                        ! See doi:10.5065/1DFH-6P97 and doi:10.1175/MWR-D-11-00215.1 for details.
            real(kind_r8), allocatable :: tv_mid_col(:) ! Virtual temperature (K) at layer midpoints of each column.
            real(kind_dyn_mpas), pointer :: rho(:, :)
            real(kind_dyn_mpas), pointer :: theta(:, :)
            real(kind_dyn_mpas), pointer :: scalars(:, :, :)

            call dyn_debug_print(debugout_verbose, 'Setting MPAS state "rho" and "theta"')

            nullify(index_qv)
            nullify(rho)
            nullify(theta)
            nullify(scalars)

            allocate(p_sfc(ncells_solve), stat=ierr)
            call check_allocate(ierr, subname, 'p_sfc(ncells_solve)', 'dyn_comp', __LINE__)

            p_sfc(:) = 0.0_kind_r8

            call dyn_set_inic_col(vc_height, lat_rad, lon_rad, global_grid_index, zint=z_int, ps=p_sfc)

            allocate(buffer_2d_real(ncells_solve, pver), stat=ierr)
            call check_allocate(ierr, subname, 'buffer_2d_real(ncells_solve, pver)', 'dyn_comp', __LINE__)

            allocate(t_mid(pver, ncells_solve), stat=ierr)
            call check_allocate(ierr, subname, 't_mid(pver, ncells_solve)', 'dyn_comp', __LINE__)

            buffer_2d_real(:, :) = 0.0_kind_r8

            call dyn_set_inic_col(vc_height, lat_rad, lon_rad, global_grid_index, zint=z_int, t=buffer_2d_real)

            ! Vertical index order is reversed between CAM-SIMA and MPAS.
            do i = 1, ncells_solve
                t_mid(:, i) = reverse(buffer_2d_real(i, :))
            end do

            deallocate(buffer_2d_real)

            allocate(p_mid_col(pver), stat=ierr)
            call check_allocate(ierr, subname, 'p_mid_col(pver)', 'dyn_comp', __LINE__)

            allocate(qv_mid_col(pver), stat=ierr)
            call check_allocate(ierr, subname, 'qv_mid_col(pver)', 'dyn_comp', __LINE__)

            allocate(tm_mid_col(pver), stat=ierr)
            call check_allocate(ierr, subname, 'tm_mid_col(pver)', 'dyn_comp', __LINE__)

            allocate(tv_mid_col(pver), stat=ierr)
            call check_allocate(ierr, subname, 'tv_mid_col(pver)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(index_qv, 'dim', 'index_qv')
            call mpas_dynamical_core % get_variable_pointer(rho, 'diag', 'rho')
            call mpas_dynamical_core % get_variable_pointer(theta, 'diag', 'theta')
            call mpas_dynamical_core % get_variable_pointer(scalars, 'state', 'scalars', time_level=1)

            ! Set `rho` and `theta` column by column. This way, peak memory usage can be reduced.
            do i = 1, ncells_solve
                qv_mid_col(:) = real(scalars(index_qv, :, i), kind_r8)
                tm_mid_col(:) = tm_of_t_qv(constant_rd, constant_rv, t_mid(:, i), qv_mid_col)
                tv_mid_col(:) = tv_of_tm_qv(tm_mid_col, qv_mid_col)

                ! Piecewise integrate hypsometric equation to derive `p_mid_col(1)`.
                ! The formulation used here is exact.
                p_mid_col(1) = p_by_hypsometric_equation( &
                    constant_g, &
                    constant_rd, &
                    p_sfc(i), &
                    real(zgrid(1, i), kind_r8), &
                    tv_mid_col(1), &
                    0.5_kind_r8 * real(zgrid(2, i) + zgrid(1, i), kind_r8))

                ! Piecewise integrate hypsometric equation to derive subsequent `p_mid_col(k)`.
                ! The formulation used here is exact.
                do k = 2, pver
                    p_mid_col(k) = p_by_hypsometric_equation( &
                        constant_g, &
                        constant_rd, &
                        p_by_hypsometric_equation( &
                            constant_g, &
                            constant_rd, &
                            p_mid_col(k - 1), &
                            0.5_kind_r8 * real(zgrid(k, i) + zgrid(k - 1, i), kind_r8), &
                            tv_mid_col(k - 1), &
                            real(zgrid(k, i), kind_r8)), &
                        real(zgrid(k, i), kind_r8), &
                        tv_mid_col(k), &
                        0.5_kind_r8 * real(zgrid(k + 1, i) + zgrid(k, i), kind_r8))
                end do

                rho(:, i) = real(rho_by_equation_of_state( &
                    constant_rd, p_mid_col, tm_mid_col), kind_dyn_mpas)
                theta(:, i) = real(theta_by_poisson_equation( &
                    constant_cpd, constant_p0, constant_rd, t_mid(:, i), p_mid_col), kind_dyn_mpas)
            end do

            deallocate(p_mid_col)
            deallocate(p_sfc)
            deallocate(qv_mid_col)
            deallocate(t_mid)
            deallocate(tm_mid_col)
            deallocate(tv_mid_col)

            nullify(index_qv)
            nullify(rho)
            nullify(theta)
            nullify(scalars)

            ! Because we are injecting data directly into MPAS memory, halo layers need to be updated manually.
            call mpas_dynamical_core % exchange_halo('rho')
            call mpas_dynamical_core % exchange_halo('theta')
        end subroutine set_mpas_state_rho_theta

        !> Set MPAS state `rho_base` (i.e., base state dry air density) and `theta_base` (i.e., base state potential temperature).
        !> (KCW, 2024-05-21)
        subroutine set_mpas_state_rho_base_theta_base()
            ! Module(s) from CAM-SIMA.
            use cam_abortutils, only: check_allocate
            use cam_logfile, only: debugout_verbose
            use dyn_grid, only: ncells_solve
            use dyn_procedure, only: p_by_hypsometric_equation, rho_by_equation_of_state, theta_by_poisson_equation
            use dynconst, only: constant_cpd => cpair, constant_g => gravit, constant_p0 => pref, &
                                constant_rd => rair
            use vert_coord, only: pver

            character(*), parameter :: subname = 'dyn_comp::set_analytic_initial_condition::set_mpas_state_rho_base_theta_base'
            integer :: i, k
            integer :: ierr
            real(kind_r8), parameter :: t_base = 250.0_kind_r8 ! Base state temperature (K) of dry isothermal atmosphere.
                                                               ! The value used here is identical to MPAS.
            real(kind_r8), allocatable :: p_base(:)            ! Base state pressure (Pa) at layer midpoints of each column.
            real(kind_dyn_mpas), pointer :: rho_base(:, :)
            real(kind_dyn_mpas), pointer :: theta_base(:, :)
            real(kind_dyn_mpas), pointer :: zz(:, :)

            call dyn_debug_print(debugout_verbose, 'Setting MPAS state "rho_base" and "theta_base"')

            nullify(rho_base)
            nullify(theta_base)
            nullify(zz)

            allocate(p_base(pver), stat=ierr)
            call check_allocate(ierr, subname, 'p_base(pver)', 'dyn_comp', __LINE__)

            call mpas_dynamical_core % get_variable_pointer(rho_base, 'diag', 'rho_base')
            call mpas_dynamical_core % get_variable_pointer(theta_base, 'diag', 'theta_base')
            call mpas_dynamical_core % get_variable_pointer(zz, 'mesh', 'zz')

            ! Set `rho_base` and `theta_base` column by column. This way, peak memory usage can be reduced.
            do i = 1, ncells_solve
                do k = 1, pver
                    ! Derive `p_base` by hypsometric equation.
                    ! The formulation used here is exact and identical to MPAS.
                    p_base(k) = p_by_hypsometric_equation( &
                        constant_g, &
                        constant_rd, &
                        constant_p0, &
                        0.0_kind_r8, &
                        t_base, &
                        0.5_kind_r8 * real(zgrid(k + 1, i) + zgrid(k, i), kind_r8))
                end do

                rho_base(:, i) = real(rho_by_equation_of_state( &
                    constant_rd, p_base, t_base) / real(zz(:, i), kind_r8), kind_dyn_mpas)
                theta_base(:, i) = real(theta_by_poisson_equation( &
                    constant_cpd, constant_p0, constant_rd, t_base, p_base), kind_dyn_mpas)
            end do

            deallocate(p_base)

            nullify(rho_base)
            nullify(theta_base)
            nullify(zz)

            ! Because we are injecting data directly into MPAS memory, halo layers need to be updated manually.
            call mpas_dynamical_core % exchange_halo('rho_base')
            call mpas_dynamical_core % exchange_halo('theta_base')
        end subroutine set_mpas_state_rho_base_theta_base
    end subroutine set_analytic_initial_condition