Make aux state BC-aware

[charleskawczynski]

Description

The update_auxiliary_state! is unaware of how boundary state is updated in boundary_state!. I think the right way to handle this to ensure that all RHS evaluations are in sync. This can be done, for example, by updating the BC non-aux variables, then call update_auxiliary_state! with the BC-aware state.

[charleskawczynski]

Once this is fixed, we should add the recover thermo state implementations back in. I've recorded what we had here, in src/Atmos/Model/thermo_states.jl:

function recover_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
)
    e_int = internal_energy(atmos, state, aux)
    return PhaseDry(atmos.param_set, e_int, state.ρ)
end

function recover_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    e_int = internal_energy(atmos, state, aux)
    return PhaseEquil{eltype(state), typeof(atmos.param_set)}(
        atmos.param_set,
        e_int,
        state.ρ,
        state.moisture.ρq_tot / state.ρ,
        aux.moisture.temperature,
    )
end

function recover_thermo_state(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::NonEquilMoist,
    state::Vars,
    aux::Vars,
)
    e_int = internal_energy(atmos, state, aux)
    q = PhasePartition(
        state.moisture.ρq_tot / state.ρ,
        state.moisture.ρq_liq / state.ρ,
        state.moisture.ρq_ice / state.ρ,
    )

    return PhaseNonEquil{eltype(state), typeof(atmos.param_set)}(
        atmos.param_set,
        e_int,
        state.ρ,
        q,
    )
end

Back into src/Atmos/Model/thermo_states.jl and

function recover_thermo_state_up(
    m::AtmosModel,
    energy::TotalEnergyModel,
    moist::Union{DryModel, EquilMoist},
    state::Vars,
    aux::Vars,
    ts::ThermodynamicState = recover_thermo_state(m, state, aux),
)
    N_up = n_updrafts(m.turbconv)
    ts_up = vuntuple(N_up) do i
        recover_thermo_state_up_i(m, state, aux, i, ts)
    end
    return ts_up
end

recover_thermo_state_up_i(m, state, aux, i_up, ts) =
    recover_thermo_state_up_i(m, m.moisture, state, aux, i_up, ts)

function recover_thermo_state_up_i(
    m::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
    i_up,
    ts::ThermodynamicState = recover_thermo_state(m, state, aux),
)
    FT = eltype(state)
    param_set = m.param_set
    up = state.turbconv.updraft

    p = air_pressure(ts)
    T_up_i = aux.turbconv.updraft[i_up].T
    q_tot_up_i = up[i_up].ρaq_tot / up[i_up].ρa
    ρ_up_i = air_density(param_set, T_up_i, p, PhasePartition(q_tot_up_i))
    q_up_i =
        PhasePartition_equil(param_set, T_up_i, ρ_up_i, q_tot_up_i, PhaseEquil)
    e_int_up_i = internal_energy(param_set, T_up_i, q_up_i)
    return PhaseEquil{FT, typeof(param_set)}(
        param_set,
        e_int_up_i,
        ρ_up_i,
        q_up_i.tot,
        T_up_i,
    )
end

function recover_thermo_state_en(
    m::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
    ts::ThermodynamicState = recover_thermo_state(m, state, aux),
)
    FT = eltype(state)
    param_set = m.param_set
    N_up = n_updrafts(m.turbconv)
    up = state.turbconv.updraft

    p = air_pressure(ts)
    T_en = aux.turbconv.environment.T
    ρ_inv = 1 / state.ρ
    q_tot = total_specific_humidity(ts)
    a_en = environment_area(state, N_up)
    q_tot_en = (q_tot - sum(vuntuple(i -> up[i].ρaq_tot, N_up)) * ρ_inv) / a_en
    ρ_en = air_density(param_set, T_en, p, PhasePartition(q_tot_en))
    q_en = PhasePartition_equil(param_set, T_en, ρ_en, q_tot_en, PhaseEquil)
    e_int_en = internal_energy(param_set, T_en, q_en)
    return PhaseEquil{FT, typeof(param_set)}(
        param_set,
        e_int_en,
        ρ_en,
        q_en.tot,
        T_en,
    )
end

into thermo_states for EDMF.