Skip to content

mpas_musica

David Fillmore edited this page Dec 12, 2025 · 1 revision

chemistry/musica/mpas_musica.F

! Copyright (c) 2025 The University Corporation for Atmospheric Research (UCAR).
!
! Unless noted otherwise source code is licensed under the BSD license.
! Additional copyright and license information can be found in the LICENSE file
! distributed with this code, or at https://mpas-dev.github.io/license.html .
!
!-----------------------------------------------------------------------
!  mpas_musica
!
!> \brief Manages interactions with the MUSICA chemistry package
!> \author CheMPAS-A Developers
!> \date   20 August 2025
!> \details
!>  This module manages the interactions with the MUSICA chemistry package,
!>  including initialization, time-stepping, and finalization for
!>  MICM (ODE solver) and TUV-x (photolysis rate constant calculator).
!--------------------------------------------------------------------------
module mpas_musica

    use musica_micm,  only: micm_t, get_micm_version
    ! get_micm_version is here to avoid an ICE when within musica_init
    use musica_state, only: state_t

    implicit none

    private

    public :: musica_init, musica_step, musica_finalize

    type(micm_t),  pointer :: micm  => null ( )  ! Pointer to the MICM ODE solver instance
    type(state_t), pointer :: state => null ( )  ! Pointer to the state of the MICM solver

    logical :: mapping_initialized = .false.
    integer, allocatable :: species_to_scalar_idx(:)
    integer :: mapped_nVertLevels = -1

    contains

    !------------------------------------------------------------------------
    !  subroutine musica_init
    !
    !> \brief Initializes the MUSICA chemistry package
    !> \author CheMPAS-A Developers
    !> \date   20 August 2025
    !> \details
    !>  This subroutine initializes the MUSICA chemistry package,
    !>  setting up necessary parameters and data structures for the simulation.
    !>  For now, we will load fixed configurations for MICM and TUV-x.
    !>  Later, we will gradually replace the fixed configuration elements
    !>  with runtime updates using the MUSICA API.
    !------------------------------------------------------------------------
    subroutine musica_init(filename_of_micm_configuration, &
            number_of_grid_cells, &
            error_code, error_message)

        use iso_fortran_env, only: real64

        use musica_micm, only : RosenbrockStandardOrder
        use musica_util, only : error_t, string_t

        use mpas_log, only : mpas_log_write

        character(len=*),              intent(in)  :: filename_of_micm_configuration
        integer,                       intent(in)  :: number_of_grid_cells
        integer,                       intent(out) :: error_code
        character(len=:), allocatable, intent(out) :: error_message

        ! TEMPORARY: ABBA specific initial concentrations
        real(real64), allocatable :: init_ab(:)
        real(real64), allocatable :: init_a(:)
        real(real64), allocatable :: init_b(:)

        character(len=*), parameter :: INITIAL_CONCENTRATION_PROPERTY = "__initial concentration"

        type(error_t)  :: error
        type(string_t) :: micm_version
        ! TEMPORARY: Hard-coded options for the MICM solver
        integer :: solver_type = RosenbrockStandardOrder
        integer :: i_species
        integer :: n_active

        micm_version = get_micm_version()

        call mpas_log_write('[MUSICA] Initializing MICM chemistry package...')
        call mpas_log_write('MICM version: ' // micm_version%value_)
        call mpas_log_write('MICM number of grid cells: $i', intArgs=[number_of_grid_cells])

        micm => micm_t(trim(filename_of_micm_configuration), solver_type, error)
        if (has_error_occurred(error, error_message, error_code)) return

        state => micm%get_state(number_of_grid_cells, error)
        if (has_error_occurred(error, error_message, error_code)) return

        do i_species = 1, state%species_ordering%size()
            call mpas_log_write('[MUSICA] MICM species: ' // state%species_ordering%name(i_species))
        end do

        mapping_initialized = .false.
        mapped_nVertLevels = -1
        if (allocated(species_to_scalar_idx)) deallocate(species_to_scalar_idx)

        ! TEMPORARY: ABBA specific initialization
        n_active = state%number_of_grid_cells
        call assign_rate_parameters(state, n_active)

        allocate(init_ab(n_active))
        allocate(init_a(n_active))
        allocate(init_b(n_active))

        init_ab = micm%get_species_property_double('AB', INITIAL_CONCENTRATION_PROPERTY, error)
        if (has_error_occurred(error, error_message, error_code)) return
        init_a = micm%get_species_property_double('A', INITIAL_CONCENTRATION_PROPERTY, error)
        if (has_error_occurred(error, error_message, error_code)) return
        init_b = micm%get_species_property_double('B', INITIAL_CONCENTRATION_PROPERTY, error)
        if (has_error_occurred(error, error_message, error_code)) return

        call set_species_profile(state, 'AB', init_ab, error)
        if (has_error_occurred(error, error_message, error_code)) return
        call set_species_profile(state, 'A', init_a, error)
        if (has_error_occurred(error, error_message, error_code)) return
        call set_species_profile(state, 'B', init_b, error)
        if (has_error_occurred(error, error_message, error_code)) return

    end subroutine musica_init

    !------------------------------------------------------------------------
    !  subroutine musica_step
    !
    !> \brief Steps the MUSICA chemistry package
    !> \author David Fillmore
    !> \date   17 November 2025
    !> \details
    !>  This subroutine steps the MICM ODE solver.
    !------------------------------------------------------------------------
    subroutine musica_step(domain, time_step, error_code, error_message)

        use iso_fortran_env, only: real64
        use mpas_kind_types, only : RKIND

        use musica_util, only : error_t
        use mpas_derived_types, only : domain_type

        use mpas_log, only : mpas_log_write

        type(domain_type), intent(inout) :: domain
        real(real64), intent(in) :: time_step

        integer,                       intent(out) :: error_code
        character(len=:), allocatable, intent(out) :: error_message

        type(error_t)        :: error
        real(kind=RKIND)     :: time_step_log

        call mpas_log_write('[MUSICA] Stepping MICM solver...')

        error_code = 0
        error_message = ''

        call MICM_from_MPAS(domain, error_code, error_message)
        if (error_code /= 0) return

        time_step_log = time_step
        call mpas_log_write('[MUSICA] MICM solve skipped (temporarily disabled) dt=$r', realArgs=[time_step_log])

        call MICM_to_MPAS(domain, error_code, error_message)
        if (error_code /= 0) return

        ! TEMPORARY: ABBA specific logging
        call mpas_log_write('[MUSICA] MICM AB concentrations')
        call log_concentrations(state, state%number_of_grid_cells, 'AB', error)
        call mpas_log_write('[MUSICA] MICM A concentrations')
        call log_concentrations(state, state%number_of_grid_cells, 'A', error)
        call mpas_log_write('[MUSICA] MICM B concentrations')
        call log_concentrations(state, state%number_of_grid_cells, 'B', error)
        call log_tracers(domain, state%number_of_grid_cells, error)
        call log_state_vs_tracer_diff(domain, 'AB', error)

    end subroutine musica_step

    !------------------------------------------------------------------------
    !  subroutine musica_finalize
    !
    !> \brief Finalizes the MUSICA chemistry package
    !> \author CheMPAS-A Developers
    !> \date   20 August 2025
    !> \details
    !>  This subroutine finalizes the MUSICA chemistry package,
    !>  cleaning up any resources and data structures used during the simulation.
    !-------------------------------------------------------------------------
    subroutine musica_finalize()

        use mpas_log, only : mpas_log_write

        call mpas_log_write('[MUSICA] Finalizing MICM chemistry package...')

    end subroutine musica_finalize

    subroutine MICM_from_MPAS(domain, error_code, error_message)
        ! Placeholder for coupling MPAS tracer state into MICM state.
        use mpas_derived_types, only : domain_type, block_type, mpas_pool_type, field3DReal
        use mpas_pool_routines, only : mpas_pool_get_subpool, mpas_pool_get_dimension, mpas_pool_get_array, mpas_pool_get_field
        use mpas_kind_types, only : RKIND
        use mpas_log, only : mpas_log_write

        type(domain_type), intent(inout) :: domain
        integer,                       intent(out) :: error_code
        character(len=:), allocatable, intent(out) :: error_message

        type(block_type), pointer :: block
        type(mpas_pool_type), pointer :: dims
        type(mpas_pool_type), pointer :: state_pool
        type(field3DReal), pointer :: scalars_field
        real(kind=RKIND), dimension(:,:,:), pointer :: scalars_tl1, scalars_tl2
        integer, pointer :: nCells_ptr, nVertLevels_ptr
        integer :: nCells, nVertLevels, num_scalars
        integer :: cell_stride, var_stride
        integer :: cell_offset, iCell, k, species_index, scalar_idx, idx
        integer :: block_id, block_cell_start, block_cell_end
        integer :: flat_idx
        integer :: idx_ab, idx_a, idx_b
        integer :: species_index_ab, species_index_a, species_index_b
        character(len=:), allocatable :: species_name, tracer_name, tracer_name_q
        character(len=:), allocatable :: scal_name
        logical :: have_mapping
    integer :: total_cells_state
    real(kind=RKIND) :: tl2_min_ab, tl2_max_ab, tl2_min_a, tl2_max_a, tl2_min_b, tl2_max_b
    real(kind=RKIND) :: state_min_ab, state_max_ab, state_min_a, state_max_a, state_min_b, state_max_b
    real(kind=RKIND) :: val_state

    error_code = 0
    error_message = ''

    if (.not. associated(state)) then
            error_code = 1
            error_message = '[MUSICA Error]: MICM state not initialized before MICM_from_MPAS'
            call mpas_log_write(error_message)
            return
        end if

        cell_stride = state%species_strides%grid_cell
        var_stride  = state%species_strides%variable
        total_cells_state = state%number_of_grid_cells

    have_mapping = mapping_initialized

    block => domain % blocklist
    cell_offset = 0
    block_id = 0
    do while (associated(block))
        block_id = block_id + 1
        nullify(dims)
        nullify(state_pool)
        dims => block % dimensions
        call mpas_pool_get_subpool(block % structs, 'state', state_pool)

            nCells = -1
            nVertLevels = -1
            nullify(nCells_ptr)
            nullify(nVertLevels_ptr)
            if (associated(dims)) then
                call mpas_pool_get_dimension(dims, 'nCells', nCells_ptr)
                call mpas_pool_get_dimension(dims, 'nVertLevels', nVertLevels_ptr)
                if (associated(nCells_ptr)) nCells = nCells_ptr
                if (associated(nVertLevels_ptr)) nVertLevels = nVertLevels_ptr
            end if

            if (.not. have_mapping) then
                nullify(scalars_field)
                call mpas_pool_get_field(state_pool, 'scalars', scalars_field, 1)
                if (.not. associated(scalars_field)) then
                    error_code = 1
                    error_message = '[MUSICA Error]: Unable to access scalars field metadata for mapping'
                    call mpas_log_write(error_message)
                    return
                end if

                num_scalars = size(scalars_field%constituentNames)
                if (.not. allocated(species_to_scalar_idx)) then
                    allocate(species_to_scalar_idx(state%species_ordering%size()))
                end if
                species_to_scalar_idx = -1

                do species_index = 1, state%species_ordering%size()
                    species_name = trim(adjustl(state%species_ordering%name(species_index)))
                    tracer_name = species_name
                    tracer_name_q = 'q' // species_name
                    do scalar_idx = 1, num_scalars
                        scal_name = trim(adjustl(scalars_field%constituentNames(scalar_idx)))
                        if (strings_match(tracer_name, scal_name) .or. strings_match(tracer_name_q, scal_name)) then
                            species_to_scalar_idx(species_index) = scalar_idx
                            exit
                        end if
                    end do
                    if (species_to_scalar_idx(species_index) < 1) then
                        error_code = 2
                        error_message = '[MUSICA Error]: No MPAS scalar found for MICM species '// &
                            trim(state%species_ordering%name(species_index))
                        call mpas_log_write(error_message)
                        return
                    else
                        call mpas_log_write('[MUSICA] Mapping MICM species '//species_name//' -> scalar index $i', &
                            intArgs=[species_to_scalar_idx(species_index)])
                    end if
                end do

                mapped_nVertLevels = nVertLevels
                mapping_initialized = .true.
                have_mapping = .true.
                call mpas_log_write('[MUSICA] MICM_from_MPAS mapping initialized for $i species', &
                    intArgs=[state%species_ordering%size()])
            else
                if (mapped_nVertLevels /= nVertLevels) then
                    error_code = 3
                    error_message = '[MUSICA Error]: Inconsistent nVertLevels across blocks for MICM mapping'
                    call mpas_log_write(error_message)
                    return
                end if
            end if

            nullify(scalars_tl1)
            nullify(scalars_tl2)
            call mpas_pool_get_array(state_pool, 'scalars', scalars_tl1, 1)
            call mpas_pool_get_array(state_pool, 'scalars', scalars_tl2, 2)
            if (.not. associated(scalars_tl1) .or. .not. associated(scalars_tl2)) then
                block => block % next
                cycle
            end if
            num_scalars = size(scalars_tl1, 1)

            idx_ab = -1
            idx_a  = -1
            idx_b  = -1
            species_index_ab = -1
            species_index_a  = -1
            species_index_b  = -1
            do species_index = 1, size(species_to_scalar_idx)
                select case (trim(adjustl(state%species_ordering%name(species_index))))
                case ('AB')
                    idx_ab = species_to_scalar_idx(species_index)
                    species_index_ab = species_index
                case ('A')
                    idx_a = species_to_scalar_idx(species_index)
                    species_index_a = species_index
                case ('B')
                    idx_b = species_to_scalar_idx(species_index)
                    species_index_b = species_index
                end select
            end do

            tl2_min_ab = huge(1.0_RKIND)
            tl2_max_ab = -tl2_min_ab
            tl2_min_a  = tl2_min_ab
            tl2_max_a  = tl2_max_ab
            tl2_min_b  = tl2_min_ab
            tl2_max_b  = tl2_max_ab
            do iCell = 1, nCells
                do k = 1, nVertLevels
                    if (idx_ab > 0) then
                        tl2_min_ab = min(tl2_min_ab, scalars_tl2(idx_ab, k, iCell))
                        tl2_max_ab = max(tl2_max_ab, scalars_tl2(idx_ab, k, iCell))
                    end if
                    if (idx_a > 0) then
                        tl2_min_a = min(tl2_min_a, scalars_tl2(idx_a, k, iCell))
                        tl2_max_a = max(tl2_max_a, scalars_tl2(idx_a, k, iCell))
                    end if
                    if (idx_b > 0) then
                        tl2_min_b = min(tl2_min_b, scalars_tl2(idx_b, k, iCell))
                        tl2_max_b = max(tl2_max_b, scalars_tl2(idx_b, k, iCell))
                    end if
                end do
            end do
            if (idx_ab < 1) then
                tl2_min_ab = 0._RKIND
                tl2_max_ab = 0._RKIND
            end if
            if (idx_a < 1) then
                tl2_min_a = 0._RKIND
                tl2_max_a = 0._RKIND
            end if
            if (idx_b < 1) then
                tl2_min_b = 0._RKIND
                tl2_max_b = 0._RKIND
            end if
            call mpas_log_write('[MUSICA] Block $i TL2 min/max qAB $r / $r qA $r / $r qB $r / $r', &
                intArgs=[block_id], realArgs=[tl2_min_ab, tl2_max_ab, tl2_min_a, tl2_max_a, tl2_min_b, tl2_max_b])

            ! Copy all cells from TL2 into TL1 so MICM reads the current state.
            do iCell = 1, nCells
                do k = 1, nVertLevels
                    do scalar_idx = 1, num_scalars
                        scalars_tl1(scalar_idx, k, iCell) = scalars_tl2(scalar_idx, k, iCell)
                    end do
                end do
            end do

            block_cell_start = cell_offset + 1
            do iCell = 1, nCells
                do k = 1, nVertLevels
                    cell_offset = cell_offset + 1
                    if (cell_offset > total_cells_state) then
                        error_code = 4
                        error_message = '[MUSICA Error]: MICM_from_MPAS exceeded MICM state size'
                        call mpas_log_write(error_message)
                        return
                    end if
                    do species_index = 1, size(species_to_scalar_idx)
                        scalar_idx = species_to_scalar_idx(species_index)
                        idx = 1 + (cell_offset - 1) * cell_stride + (species_index - 1) * var_stride
                        state%concentrations(idx) = scalars_tl1(scalar_idx, k, iCell)
                    end do
                end do
            end do
            block_cell_end = cell_offset

            state_min_ab = huge(1.0_RKIND)
            state_max_ab = -state_min_ab
            state_min_a  = state_min_ab
            state_max_a  = state_max_ab
            state_min_b  = state_min_ab
            state_max_b  = state_max_ab
            if (species_index_ab > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_ab - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_ab = min(state_min_ab, val_state)
                    state_max_ab = max(state_max_ab, val_state)
                end do
            end if
            if (species_index_a > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_a - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_a = min(state_min_a, val_state)
                    state_max_a = max(state_max_a, val_state)
                end do
            end if
            if (species_index_b > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_b - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_b = min(state_min_b, val_state)
                    state_max_b = max(state_max_b, val_state)
                end do
            end if
            if (species_index_ab < 1) then
                state_min_ab = 0._RKIND
                state_max_ab = 0._RKIND
            end if
            if (species_index_a < 1) then
                state_min_a = 0._RKIND
                state_max_a = 0._RKIND
            end if
            if (species_index_b < 1) then
                state_min_b = 0._RKIND
                state_max_b = 0._RKIND
            end if
            call mpas_log_write('[MUSICA] Block $i MICM state min/max AB $r / $r A $r / $r B $r / $r', &
                intArgs=[block_id], realArgs=[state_min_ab, state_max_ab, state_min_a, state_max_a, state_min_b, state_max_b])

            ! Sample the first cell/level mapping to check index math.
            if (block_id == 1 .and. nCells > 0 .and. nVertLevels > 0) then
                flat_idx = block_cell_start + (1 - 1) * nVertLevels + 1 - 1
                if (idx_ab > 0 .and. species_index_ab > 0) then
                    idx = 1 + flat_idx * cell_stride + (species_index_ab - 1) * var_stride
                    call mpas_log_write('[MUSICA] Sample C1K1 qAB TL2=$r MICM=$r', &
                        realArgs=[scalars_tl2(idx_ab, 1, 1), real(state%concentrations(idx), kind=RKIND)])
                end if
                if (idx_a > 0 .and. species_index_a > 0) then
                    idx = 1 + flat_idx * cell_stride + (species_index_a - 1) * var_stride
                    call mpas_log_write('[MUSICA] Sample C1K1 qA  TL2=$r MICM=$r', &
                        realArgs=[scalars_tl2(idx_a, 1, 1), real(state%concentrations(idx), kind=RKIND)])
                end if
                if (idx_b > 0 .and. species_index_b > 0) then
                    idx = 1 + flat_idx * cell_stride + (species_index_b - 1) * var_stride
                    call mpas_log_write('[MUSICA] Sample C1K1 qB  TL2=$r MICM=$r', &
                        realArgs=[scalars_tl2(idx_b, 1, 1), real(state%concentrations(idx), kind=RKIND)])
                end if
            end if

            block => block % next
        end do

        if (cell_offset /= total_cells_state) then
            error_code = 5
            error_message = '[MUSICA Error]: MICM_from_MPAS grid-cell count mismatch with MICM state'
            call mpas_log_write('[MUSICA Error]: MICM_from_MPAS filled $i of $i grid cells (mismatch with MICM state)', &
                intArgs=[cell_offset, total_cells_state])
            return
        end if
    end subroutine MICM_from_MPAS

    subroutine MICM_to_MPAS(domain, error_code, error_message)
        ! Placeholder for coupling MICM state back into MPAS tracers.
        use mpas_derived_types, only : domain_type, block_type, mpas_pool_type
        use mpas_pool_routines, only : mpas_pool_get_subpool, mpas_pool_get_dimension, mpas_pool_get_array
        use mpas_kind_types, only : RKIND
        use mpas_log, only : mpas_log_write

        type(domain_type), intent(inout) :: domain
        integer,                       intent(out) :: error_code
        character(len=:), allocatable, intent(out) :: error_message

        type(block_type), pointer :: block
        type(mpas_pool_type), pointer :: dims
        type(mpas_pool_type), pointer :: state_pool
        real(kind=RKIND), dimension(:,:,:), pointer :: scalars_tl1, scalars_tl2
        integer, pointer :: nCells_ptr, nVertLevels_ptr
        integer :: nCells, nVertLevels
        integer :: cell_stride, var_stride
        integer :: cell_offset, iCell, k, species_index, scalar_idx, idx
        integer :: block_id, block_cell_start, block_cell_end
        integer :: idx_ab, idx_a, idx_b
        integer :: species_index_ab, species_index_a, species_index_b
        integer :: total_cells_state
        real(kind=RKIND) :: state_min_ab, state_max_ab, state_min_a, state_max_a, state_min_b, state_max_b
        real(kind=RKIND) :: tl2_min_ab, tl2_max_ab, tl2_min_a, tl2_max_a, tl2_min_b, tl2_max_b
        real(kind=RKIND) :: val_state

        error_code = 0
        error_message = ''

        if (.not. associated(state)) then
            error_code = 1
            error_message = '[MUSICA Error]: MICM state not initialized before MICM_to_MPAS'
            call mpas_log_write(error_message)
            return
        end if

        if (.not. mapping_initialized) then
            call mpas_log_write('[MUSICA] MICM_to_MPAS: mapping not initialized, skipping update')
            return
        end if

        cell_stride = state%species_strides%grid_cell
        var_stride  = state%species_strides%variable
        total_cells_state = state%number_of_grid_cells

        block => domain % blocklist
        cell_offset = 0
        block_id = 0
        do while (associated(block))
            block_id = block_id + 1
            nullify(dims)
            nullify(state_pool)
            dims => block % dimensions
            call mpas_pool_get_subpool(block % structs, 'state', state_pool)

            nCells = -1
            nVertLevels = -1
            nullify(nCells_ptr)
            nullify(nVertLevels_ptr)
            if (associated(dims)) then
                call mpas_pool_get_dimension(dims, 'nCells', nCells_ptr)
                call mpas_pool_get_dimension(dims, 'nVertLevels', nVertLevels_ptr)
                if (associated(nCells_ptr)) nCells = nCells_ptr
                if (associated(nVertLevels_ptr)) nVertLevels = nVertLevels_ptr
            end if

            if (mapped_nVertLevels /= nVertLevels) then
                error_code = 2
                error_message = '[MUSICA Error]: Inconsistent nVertLevels during MICM_to_MPAS'
                call mpas_log_write(error_message)
                return
            end if

            nullify(scalars_tl1)
            nullify(scalars_tl2)
            call mpas_pool_get_array(state_pool, 'scalars', scalars_tl1, 1)
            call mpas_pool_get_array(state_pool, 'scalars', scalars_tl2, 2)

            idx_ab = -1
            idx_a  = -1
            idx_b  = -1
            species_index_ab = -1
            species_index_a  = -1
            species_index_b  = -1
            do species_index = 1, size(species_to_scalar_idx)
                select case (trim(adjustl(state%species_ordering%name(species_index))))
                case ('AB')
                    idx_ab = species_to_scalar_idx(species_index)
                    species_index_ab = species_index
                case ('A')
                    idx_a = species_to_scalar_idx(species_index)
                    species_index_a = species_index
                case ('B')
                    idx_b = species_to_scalar_idx(species_index)
                    species_index_b = species_index
                end select
            end do

            block_cell_start = cell_offset + 1
            block_cell_end = block_cell_start + (nCells * nVertLevels) - 1
            if (block_cell_end > total_cells_state) then
                error_code = 3
                error_message = '[MUSICA Error]: MICM_to_MPAS exceeded MICM state size'
                call mpas_log_write(error_message)
                return
            end if

            state_min_ab = huge(1.0_RKIND)
            state_max_ab = -state_min_ab
            state_min_a  = state_min_ab
            state_max_a  = state_max_ab
            state_min_b  = state_min_ab
            state_max_b  = state_max_b
            if (species_index_ab > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_ab - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_ab = min(state_min_ab, val_state)
                    state_max_ab = max(state_max_ab, val_state)
                end do
            end if
            if (species_index_a > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_a - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_a = min(state_min_a, val_state)
                    state_max_a = max(state_max_a, val_state)
                end do
            end if
            if (species_index_b > 0) then
                do idx = block_cell_start, block_cell_end
                    scalar_idx = 1 + (idx - 1) * cell_stride + (species_index_b - 1) * var_stride
                    val_state = real(state%concentrations(scalar_idx), kind=RKIND)
                    state_min_b = min(state_min_b, val_state)
                    state_max_b = max(state_max_b, val_state)
                end do
            end if
            if (species_index_ab < 1) then
                state_min_ab = 0._RKIND
                state_max_ab = 0._RKIND
            end if
            if (species_index_a < 1) then
                state_min_a = 0._RKIND
                state_max_a = 0._RKIND
            end if
            if (species_index_b < 1) then
                state_min_b = 0._RKIND
                state_max_b = 0._RKIND
            end if
            call mpas_log_write('[MUSICA] Block $i MICM state before write AB $r / $r A $r / $r B $r / $r', &
                intArgs=[block_id], realArgs=[state_min_ab, state_max_ab, state_min_a, state_max_a, state_min_b, state_max_b])

            do iCell = 1, nCells
                do k = 1, nVertLevels
                    cell_offset = cell_offset + 1
                    do species_index = 1, size(species_to_scalar_idx)
                        scalar_idx = species_to_scalar_idx(species_index)
                        idx = 1 + (cell_offset - 1) * cell_stride + (species_index - 1) * var_stride
                        scalars_tl2(scalar_idx, k, iCell) = state%concentrations(idx)
                    end do
                end do
            end do
            block_cell_end = cell_offset

            tl2_min_ab = huge(1.0_RKIND)
            tl2_max_ab = -tl2_min_ab
            tl2_min_a  = tl2_min_ab
            tl2_max_a  = tl2_max_ab
            tl2_min_b  = tl2_min_ab
            tl2_max_b  = tl2_max_ab
            do iCell = 1, nCells
                do k = 1, nVertLevels
                    if (idx_ab > 0) then
                        tl2_min_ab = min(tl2_min_ab, scalars_tl2(idx_ab, k, iCell))
                        tl2_max_ab = max(tl2_max_ab, scalars_tl2(idx_ab, k, iCell))
                    end if
                    if (idx_a > 0) then
                        tl2_min_a = min(tl2_min_a, scalars_tl2(idx_a, k, iCell))
                        tl2_max_a = max(tl2_max_a, scalars_tl2(idx_a, k, iCell))
                    end if
                    if (idx_b > 0) then
                        tl2_min_b = min(tl2_min_b, scalars_tl2(idx_b, k, iCell))
                        tl2_max_b = max(tl2_max_b, scalars_tl2(idx_b, k, iCell))
                    end if
                end do
            end do
            if (idx_ab < 1) then
                tl2_min_ab = 0._RKIND
                tl2_max_ab = 0._RKIND
            end if
            if (idx_a < 1) then
                tl2_min_a = 0._RKIND
                tl2_max_a = 0._RKIND
            end if
            if (idx_b < 1) then
                tl2_min_b = 0._RKIND
                tl2_max_b = 0._RKIND
            end if
            call mpas_log_write('[MUSICA] Block $i TL2 after MICM write qAB $r / $r qA $r / $r qB $r / $r', &
                intArgs=[block_id], realArgs=[tl2_min_ab, tl2_max_ab, tl2_min_a, tl2_max_a, tl2_min_b, tl2_max_b])

            block => block % next
        end do

        if (cell_offset /= total_cells_state) then
            error_code = 4
            error_message = '[MUSICA Error]: MICM_to_MPAS grid-cell count mismatch with MICM state'
            call mpas_log_write('[MUSICA Error]: MICM_to_MPAS updated $i of $i grid cells (mismatch with MICM state)', &
                intArgs=[cell_offset, total_cells_state])
        end if
    end subroutine MICM_to_MPAS

    subroutine assign_rate_parameters(state, number_of_grid_cells)
        ! Provide a default value for any rate parameters (none for ABBA by default).
        use iso_fortran_env, only : real64

        type(state_t), intent(inout) :: state
        integer, intent(in)          :: number_of_grid_cells

        integer :: cell_stride, var_stride, rp_index, cell, idx

        if (state%number_of_rate_parameters == 0) return

        cell_stride = state%rate_parameters_strides%grid_cell
        var_stride  = state%rate_parameters_strides%variable

        do rp_index = 1, state%rate_parameters_ordering%size()
            do cell = 1, number_of_grid_cells
                idx = 1 + (cell - 1) * cell_stride + (rp_index - 1) * var_stride
                state%rate_parameters(idx) = 1.0_real64
            end do
        end do
    end subroutine assign_rate_parameters

    subroutine set_species_profile(state, species_name, values, err)
        ! Write a 1-D profile into the contiguous concentrations array.
        use iso_fortran_env, only : real64
        use musica_util, only : error_t, string_t
        use mpas_log, only : mpas_log_write

        type(state_t), intent(inout) :: state
        character(len=*), intent(in) :: species_name
        real(real64), intent(in)     :: values(:)
        type(error_t), intent(inout) :: err
        integer :: species_index, cell_stride, var_stride, cell, idx

        species_index = find_species_index_by_name(state, species_name)
        if (species_index < 1) then
            call mpas_log_write('[MUSICA] set_species_profile: species '//trim(species_name)//' not found')
            return
        end if

        cell_stride = state%species_strides%grid_cell
        var_stride  = state%species_strides%variable

        do cell = 1, size(values)
            idx = 1 + (cell - 1) * cell_stride + (species_index - 1) * var_stride
            state%concentrations(idx) = values(cell)
        end do
    end subroutine set_species_profile

    subroutine log_concentrations(state, number_of_grid_cells, species_name, err)
        use iso_fortran_env, only : real64
        use mpas_log, only        : mpas_log_write
        use mpas_kind_types, only : RKIND
        use musica_util, only : error_t, string_t

        type(state_t), intent(in)    :: state
        integer, intent(in)          :: number_of_grid_cells
        character(len=*), intent(in) :: species_name
        type(error_t), intent(inout) :: err

        integer :: species_index, cell_stride, var_stride
        integer :: start_idx, end_idx, sample_count, flat_idx, slice_len
        real(real64), pointer :: species_slice(:)
        real (kind=RKIND) :: concentration

        species_index = find_species_index_by_name(state, species_name)
        if (species_index < 1) then
            call mpas_log_write('[MUSICA] log_concentrations: species '//trim(species_name)//' not found')
            return
        end if

        cell_stride = state%species_strides%grid_cell
        var_stride  = state%species_strides%variable

        if (cell_stride < 1 .or. var_stride < 1) then
            call mpas_log_write('[MUSICA] log_concentrations: invalid strides for '//trim(species_name)// &
                ' (cell_stride=$i, var_stride=$i)', intArgs=[cell_stride, var_stride])
            return
        end if

        start_idx = 1 + (species_index - 1) * var_stride
        end_idx   = start_idx + (number_of_grid_cells - 1) * cell_stride

        if (start_idx < 1 .or. start_idx > size(state%concentrations)) then
            call mpas_log_write('[MUSICA] log_concentrations: start index $i out of bounds (size=$i) for '// &
                trim(species_name), intArgs=[start_idx, size(state%concentrations)])
            return
        end if

        end_idx = min(end_idx, size(state%concentrations))
        sample_count = 1 + (end_idx - start_idx) / cell_stride
        if (sample_count < 1) then
            call mpas_log_write('[MUSICA] log_concentrations: empty slice for '//trim(species_name))
            return
        end if

        species_slice => state%concentrations(start_idx:end_idx:cell_stride)
        slice_len = size(species_slice)
        sample_count = min(slice_len, 10)

        call mpas_log_write('[MUSICA] log_concentrations '//trim(species_name)// &
            ' idx=$i start=$i stride=$i var_stride=$i samples=$i', &
            intArgs=[species_index, start_idx, cell_stride, var_stride, sample_count])

        do flat_idx = 1, sample_count
            concentration = species_slice(flat_idx)
            call mpas_log_write('  MICM[$i] idx=$i conc=$r', &
                intArgs=[flat_idx, start_idx + (flat_idx - 1) * cell_stride], realArgs=[concentration])
        end do
    end subroutine log_concentrations

    integer function find_species_index_by_name(state, species_name)
        ! Return 1-based index of species_name within state%species_ordering%name(:).
        type(state_t), intent(in)    :: state
        character(len=*), intent(in) :: species_name
        integer :: i

        find_species_index_by_name = -1
        do i = 1, state%species_ordering%size()
            if (strings_match(trim(adjustl(state%species_ordering%name(i))), trim(adjustl(species_name)))) then
                find_species_index_by_name = i
                return
            end if
        end do
    end function find_species_index_by_name

    subroutine log_solver_stats(solver_stats)
        use mpas_log, only        : mpas_log_write
        use mpas_kind_types, only : RKIND
        use musica_micm, only     : solver_stats_t

        type(solver_stats_t), intent(in) :: solver_stats
 
        integer :: function_calls, jacobian_updates, number_of_steps
        integer :: accepted, rejected, decompositions, solves

        real (kind=RKIND) :: final_time

        function_calls = solver_stats%function_calls()
        jacobian_updates = solver_stats%jacobian_updates()
        number_of_steps = solver_stats%number_of_steps()
        accepted = solver_stats%accepted()
        rejected = solver_stats%rejected()
        decompositions = solver_stats%decompositions()
        solves = solver_stats%solves()
        final_time = solver_stats%final_time()

        call mpas_log_write('[MUSICA] MICM Solver statistics ...')
        call mpas_log_write('  MICM function calls: $i', intArgs=[function_calls])
        call mpas_log_write('  MICM jacobian updates: $i', intArgs=[jacobian_updates])
        call mpas_log_write('  MICM number of steps: $i', intArgs=[number_of_steps])
        call mpas_log_write('  MICM accepted: $i', intArgs=[accepted])
        call mpas_log_write('  MICM rejected: $i', intArgs=[rejected])
        call mpas_log_write('  MICM LU decompositions: $i', intArgs=[decompositions])
        call mpas_log_write('  MICM linear solves: $i', intArgs=[solves])
        call mpas_log_write('  MICM final time (s): $r', realArgs=[final_time])
    end subroutine log_solver_stats

    subroutine log_tracers(domain, limit_cells, err)
        ! Log MPAS tracer values for AB, A, B alongside MICM concentrations.
        use mpas_derived_types, only : domain_type, block_type, mpas_pool_type
        use mpas_pool_routines, only : mpas_pool_get_subpool, mpas_pool_get_array
        use mpas_log, only : mpas_log_write
        use mpas_kind_types, only : RKIND
        use musica_util, only : error_t

        type(domain_type), intent(inout) :: domain
        integer, intent(in)              :: limit_cells
        type(error_t), intent(inout)     :: err

        type(block_type), pointer :: block
        type(mpas_pool_type), pointer :: state_pool
        real(kind=RKIND), dimension(:,:,:), pointer :: scalars
        integer :: nCells, nVertLevels_block, k, iCell
        integer :: idx_ab, idx_a, idx_b
        integer :: flat_counter, flat_mid, flat_start, flat_end
        integer :: nVertLevels_ref, mid_level
        integer :: global_cell
        logical :: range_set

        if (.not. mapping_initialized) then
            call mpas_log_write('[MUSICA] log_tracers skipped: mapping not initialized')
            return
        end if

        block => domain % blocklist
        flat_counter = 0
        range_set = .false.
        nVertLevels_ref = -1
        do while (associated(block) .and. flat_counter < limit_cells)
            nullify(state_pool)
            call mpas_pool_get_subpool(block % structs, 'state', state_pool)
            if (.not. associated(state_pool)) then
                block => block % next
                cycle
            end if

            nullify(scalars)
            call mpas_pool_get_array(state_pool, 'scalars', scalars, 1)
            if (.not. associated(scalars)) then
                block => block % next
                cycle
            end if

            nCells = size(scalars, 3)
            nVertLevels_block = size(scalars, 2)
            if (nVertLevels_ref < 0) then
                nVertLevels_ref = nVertLevels_block
                mid_level = max(1, min(nVertLevels_ref, (nVertLevels_ref + 1) / 2))
            else if (nVertLevels_block /= nVertLevels_ref) then
                call mpas_log_write('[MUSICA] log_tracers skipped: inconsistent nVertLevels across blocks')
                return
            end if

            if (.not. range_set .and. nVertLevels_ref > 0) then
                flat_mid = max(1, min(limit_cells, (limit_cells + 1) / 2))
                flat_start = max(1, flat_mid - 4)
                flat_end = min(limit_cells, flat_start + 9)
                flat_start = max(1, flat_end - 9)
                range_set = .true.
            end if

            idx_ab = species_to_scalar_idx(find_species_index_by_name(state, 'AB'))
            idx_a  = species_to_scalar_idx(find_species_index_by_name(state, 'A'))
            idx_b  = species_to_scalar_idx(find_species_index_by_name(state, 'B'))
            do iCell = 1, nCells
                do k = 1, nVertLevels_block
                    flat_counter = flat_counter + 1
                    if (.not. range_set) cycle
                    if (flat_counter < flat_start) cycle
                    if (flat_counter > flat_end) exit
                    if (k /= mid_level) cycle
                    global_cell = ((flat_counter - mid_level) / nVertLevels_ref) + 1
                    call mpas_log_write('[MUSICA] Tracers cell=$i level=$i qAB=$r qA=$r qB=$r', &
                        intArgs=[global_cell, k], &
                        realArgs=[scalars(idx_ab, k, iCell), scalars(idx_a, k, iCell), scalars(idx_b, k, iCell)])
                end do
                if (flat_counter > flat_end) exit
            end do

            block => block % next
            if (range_set .and. flat_counter >= flat_end) exit
        end do
    end subroutine log_tracers

    subroutine log_state_vs_tracer_diff(domain, species_name, err)
        ! Compute a max-norm difference between MICM state and MPAS TL2 tracers for one species.
        use mpas_derived_types, only : domain_type, block_type, mpas_pool_type
        use mpas_pool_routines, only : mpas_pool_get_subpool, mpas_pool_get_array
        use mpas_log, only : mpas_log_write
        use mpas_kind_types, only : RKIND
        use musica_util, only : error_t

        type(domain_type), intent(inout) :: domain
        character(len=*), intent(in)    :: species_name
        type(error_t), intent(inout)    :: err

        type(block_type), pointer :: block
        type(mpas_pool_type), pointer :: state_pool
        real(kind=RKIND), dimension(:,:,:), pointer :: scalars_tl2
        integer :: species_index, scalar_idx
        integer :: cell_stride, var_stride
        integer :: nCells, nVertLevels, iCell, k
        integer :: total_cells_state, cell_offset
        integer :: block_id, max_block, max_flat_cell, max_level, idx_state_max, idx_state
        real(kind=RKIND) :: diff, max_diff
        logical :: reached_end

        if (.not. mapping_initialized) then
            call mpas_log_write('[MUSICA] log_state_vs_tracer_diff skipped: mapping not initialized')
            return
        end if

        species_index = find_species_index_by_name(state, species_name)
        if (species_index < 1) then
            call mpas_log_write('[MUSICA] log_state_vs_tracer_diff: species '//trim(species_name)//' not found')
            return
        end if

        scalar_idx = species_to_scalar_idx(species_index)
        if (scalar_idx < 1) then
            call mpas_log_write('[MUSICA] log_state_vs_tracer_diff: no scalar index for species '//trim(species_name))
            return
        end if

        cell_stride = state%species_strides%grid_cell
        var_stride  = state%species_strides%variable
        total_cells_state = state%number_of_grid_cells

        block => domain % blocklist
        cell_offset = 0
        block_id = 0
        max_block = -1
        max_flat_cell = -1
        max_level = -1
        idx_state_max = -1
        max_diff = -1._RKIND
        reached_end = .false.

        do while (associated(block) .and. .not. reached_end)
            block_id = block_id + 1
            nullify(state_pool)
            call mpas_pool_get_subpool(block % structs, 'state', state_pool)
            if (.not. associated(state_pool)) then
                block => block % next
                cycle
            end if

            nullify(scalars_tl2)
            call mpas_pool_get_array(state_pool, 'scalars', scalars_tl2, 2)
            if (.not. associated(scalars_tl2)) then
                block => block % next
                cycle
            end if

            nCells = size(scalars_tl2, 3)
            nVertLevels = size(scalars_tl2, 2)

            do iCell = 1, nCells
                do k = 1, nVertLevels
                    if (cell_offset >= total_cells_state) then
                        reached_end = .true.
                        exit
                    end if
                    cell_offset = cell_offset + 1
                    idx_state = 1 + (cell_offset - 1) * cell_stride + (species_index - 1) * var_stride
                    diff = abs(state%concentrations(idx_state) - scalars_tl2(scalar_idx, k, iCell))
                    if (diff > max_diff) then
                        max_diff = diff
                        max_block = block_id
                        max_flat_cell = cell_offset
                        max_level = k
                        idx_state_max = idx_state
                    end if
                end do
            end do

            block => block % next
        end do

        if (cell_offset /= total_cells_state) then
            call mpas_log_write('[MUSICA] log_state_vs_tracer_diff filled $i of $i cells (mismatch)', &
                intArgs=[cell_offset, total_cells_state])
        end if

        if (max_diff >= 0._RKIND) then
            call mpas_log_write('[MUSICA] Consistency q'//trim(species_name)// &
                ' max|MICM - TL2|=$r at flat_cell=$i level=$i block=$i state_idx=$i', &
                realArgs=[max_diff], intArgs=[max_flat_cell, max_level, max_block, idx_state_max])
        else
            call mpas_log_write('[MUSICA] Consistency q'//trim(species_name)//' no samples found')
        end if
    end subroutine log_state_vs_tracer_diff

    !-------------------------------------------------------------------------
    !  function has_error_occurred
    !
    !> \author CheMPAS-A Developers
    !> \date   20 August 2025
    !  \details
    !>  Evaluate a MUSICA error for failure and convert to error data
    !>  @param[in] error The error code to evaluate and convert.
    !>  @param[out] error_message The error message.
    !>  @param[out] error_code The error code.
    !>  @return True for an error, false for success.
    !-------------------------------------------------------------------------
    logical function has_error_occurred(error, error_message, error_code)
        use musica_util, only: error_t

        type(error_t),                 intent(in)  :: error
        character(len=:), allocatable, intent(out) :: error_message
        integer,                       intent(out) :: error_code

        character(len=30) :: error_code_str

        if ( error%is_success( ) ) then
          error_code = 0
          error_message = ''
          has_error_occurred = .false.
          return
        end if
        error_code = error%code( )
        write(error_code_str, '(I30)') error%code( )
        error_message = '[MUSICA Error]: ' // error%category( ) // '[' // &
                        trim( adjustl( error_code_str ) ) // ']: ' // error%message( )
        has_error_occurred = .true.
    end function has_error_occurred

    logical function strings_match(a, b)
        ! Case-insensitive comparison of two trimmed strings.
        character(len=*), intent(in) :: a, b
        integer :: i, la, lb
        la = len_trim(a)
        lb = len_trim(b)
        if (la /= lb) then
            strings_match = .false.
            return
        end if
        strings_match = .true.
        do i = 1, la
            strings_match = strings_match .and. to_lower_char(a(i:i)) == to_lower_char(b(i:i))
        end do
    end function strings_match

    character(len=1) function to_lower_char(ch)
        ! Convert A-Z to a-z; leave other characters unchanged.
        character(len=*), intent(in) :: ch
        integer :: ia
        ia = iachar(ch)
        if (ia >= iachar('A') .and. ia <= iachar('Z')) then
            to_lower_char = achar(ia + 32)
        else
            to_lower_char = ch
        end if
    end function to_lower_char

end module mpas_musica

Clone this wiki locally