Skip to content

Commit

Permalink
restructure component models
Browse files Browse the repository at this point in the history
  • Loading branch information
juliasloan25 committed Mar 7, 2024
1 parent 363305e commit 18856c1
Show file tree
Hide file tree
Showing 13 changed files with 654 additions and 620 deletions.
2 changes: 1 addition & 1 deletion docs/src/interfacer.md
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ information needed to run that simulation.

Each `ComponentModelSimulation` must extend the following functions to be able
to use our coupler. For some existing models, these are defined within
ClimaCoupler.jl in that model’s `init.jl` file, but it is preferable
ClimaCoupler.jl in that model’s file in `experiments/AMIP/components/`, but it is preferable
for these to be defined in a model’s own repository. Note that the dispatch
`::ComponentModelSimulation` in the function definitions given below should
be replaced with the particular component model extending these functions.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ import ClimaCoupler.Utilities: swap_space!

include("climaatmos_extra_diags.jl")

# the clima atmos `integrator` is now defined
###
### Functions required by ClimaCoupler.jl for an AtmosModelSimulation
###
struct ClimaAtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
params::P
Y_init::Y
Expand All @@ -31,43 +33,6 @@ struct ClimaAtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
end
name(::ClimaAtmosSimulation) = "ClimaAtmosSimulation"

"""
get_atmos_config(coupler_dict::Dict)
Returns the specified atmospheric configuration (`atmos_config_dict`) overwitten by arguments
in the coupler dictionary (`config_dict`).
"""
function get_atmos_config(coupler_dict)
atmos_config_file = coupler_dict["atmos_config_file"]
# override default or specified configs with coupler arguments, and set the correct atmos config_file
if isnothing(atmos_config_file)
@info "Using Atmos default configuration"
atmos_config = merge(CA.default_config_dict(), coupler_dict, Dict("config_file" => atmos_config_file))
else
@info "Using Atmos configuration from $atmos_config_file"
atmos_config = merge(
CA.override_default_config(joinpath(pkgdir(CA), atmos_config_file)),
coupler_dict,
Dict("config_file" => atmos_config_file),
)
end

# use coupler toml if atmos is not defined
atmos_toml_file = atmos_config["toml"]
coupler_toml_file = coupler_dict["coupler_toml_file"]
default_toml_file = "toml/default_coarse.toml"

toml_file = !isempty(atmos_toml_file) ? joinpath(pkgdir(CA), atmos_toml_file[1]) : nothing
toml_file = !isnothing(coupler_toml_file) ? joinpath(pkgdir(ClimaCoupler), coupler_toml_file) : toml_file
toml_file = isnothing(toml_file) ? joinpath(pkgdir(ClimaCoupler), default_toml_file) : toml_file

if !isnothing(toml_file)
@info "Overwriting Atmos parameters from $toml_file"
atmos_config = merge(atmos_config, Dict("toml" => [toml_file]))
end
return atmos_config
end

function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
# By passing `parsed_args` to `AtmosConfig`, `parsed_args` overwrites the default atmos config
atmos_config = CA.AtmosConfig(atmos_config_dict)
Expand Down Expand Up @@ -110,17 +75,78 @@ function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
return sim
end

"""
get_model_prog_state(sim::ClimaAtmosSimulation)
Extension of Checkpointer.get_model_prog_state to get the model state.
"""
function get_model_prog_state(sim::ClimaAtmosSimulation)
return sim.integrator.u
end

"""
get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
Extension of Interfacer.get_field to get the net TOA radiation, which is a sum of the
upward and downward longwave and shortwave radiation.
"""
function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
FT = eltype(atmos_sim.integrator.u)
# save radiation source
if atmos_sim.integrator.p.radiation.radiation_model != nothing
face_space = axes(atmos_sim.integrator.u.f)
nz_faces = length(Spaces.vertical_topology(face_space).mesh.faces)

(; face_lw_flux_dn, face_lw_flux_up, face_sw_flux_dn, face_sw_flux_up) =
atmos_sim.integrator.p.radiation.radiation_model

LWd_TOA = Fields.level(CA.RRTMGPI.array2field(FT.(face_lw_flux_dn), face_space), nz_faces - half)
LWu_TOA = Fields.level(CA.RRTMGPI.array2field(FT.(face_lw_flux_up), face_space), nz_faces - half)
SWd_TOA = Fields.level(CA.RRTMGPI.array2field(FT.(face_sw_flux_dn), face_space), nz_faces - half)
SWu_TOA = Fields.level(CA.RRTMGPI.array2field(FT.(face_sw_flux_up), face_space), nz_faces - half)

return @. -(LWd_TOA + SWd_TOA - LWu_TOA - SWu_TOA) # [W/m^2]
else
return FT(0)
end
end

function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy})
thermo_params = get_thermo_params(atmos_sim)

ᶜS_ρq_tot = atmos_sim.integrator.p.precipitation.ᶜS_ρq_tot
ᶜts = atmos_sim.integrator.p.precomputed.ᶜts
ᶜΦ = atmos_sim.integrator.p.core.ᶜΦ

# return total energy and (if Microphysics0Moment) the energy lost due to precipitation removal
if atmos_sim.integrator.p.atmos.precip_model isa CA.Microphysics0Moment
ᶜS_ρq_tot = atmos_sim.integrator.p.precipitation.ᶜS_ρq_tot
ᶜts = atmos_sim.integrator.p.precomputed.ᶜts
ᶜΦ = atmos_sim.integrator.p.core.ᶜΦ
return atmos_sim.integrator.u.c.ρe_tot .-
ᶜS_ρq_tot .* CA.e_tot_0M_precipitation_sources_helper.(Ref(thermo_params), ᶜts, ᶜΦ) .*
atmos_sim.integrator.dt
else
return atmos_sim.integrator.u.c.ρe_tot
end
end

# extensions required by the Interfacer
get_field(sim::ClimaAtmosSimulation, ::Val{:air_density}) =
TD.air_density.(thermo_params, sim.integrator.p.precomputed.ᶜts)
get_field(sim::ClimaAtmosSimulation, ::Val{:air_temperature}) =
TD.air_temperature.(thermo_params, sim.integrator.p.precomputed.ᶜts)
get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) = sim.integrator.p.precipitation.col_integrated_rain # kg/m^2/s
get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_sfc}) =
ClimaCore.Fields.level(sim.integrator.p.radiation.ᶠradiation_flux, half)
get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) = sim.integrator.p.precipitation.col_integrated_rain # kg/m^2/s
get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) = sim.integrator.p.precipitation.col_integrated_snow # kg/m^2/s
get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
ClimaCore.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_h_tot)
get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
ClimaCore.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot)
get_field(sim::ClimaAtmosSimulation, ::Val{:thermo_state_int}) =
ClimaCore.Spaces.level(sim.integrator.p.precomputed.ᶜts, 1)
get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot

# extensions required by FluxCalculator (partitioned fluxes)
get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
Expand All @@ -131,12 +157,6 @@ function get_field(sim::ClimaAtmosSimulation, ::Val{:uv_int})
uₕ_int = ClimaCore.Geometry.UVVector.(ClimaCore.Spaces.level(sim.integrator.u.c.uₕ, 1))
return @. StaticArrays.SVector(uₕ_int.components.data.:1, uₕ_int.components.data.:2)
end
get_field(sim::ClimaAtmosSimulation, ::Val{:air_density}) =
TD.air_density.(thermo_params, sim.integrator.p.precomputed.ᶜts)
get_field(sim::ClimaAtmosSimulation, ::Val{:air_temperature}) =
TD.air_temperature.(thermo_params, sim.integrator.p.precomputed.ᶜts)

get_surface_params(sim::ClimaAtmosSimulation) = CAP.surface_fluxes_params(sim.integrator.p.params)

function update_field!(atmos_sim::ClimaAtmosSimulation, ::Val{:co2}, field)
if atmos_sim.integrator.p.atmos.radiation_mode isa CA.RRTMGPI.GrayRadiation
Expand All @@ -158,7 +178,6 @@ function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_albedo}, field)
reshape(CA.RRTMGPI.field2array(field), 1, length(parent(field)))
end

# get_surface_params required by FluxCalculator (partitioned fluxes)
function update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fields)
(; F_turb_energy, F_turb_moisture, F_turb_ρτxz, F_turb_ρτyz) = fields

Expand Down Expand Up @@ -188,7 +207,6 @@ end

# extensions required by FieldExchanger
step!(sim::ClimaAtmosSimulation, t) = step!(sim.integrator, t - sim.integrator.t, true)

reinit!(sim::ClimaAtmosSimulation) = reinit!(sim.integrator)

function update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
Expand All @@ -200,6 +218,49 @@ function update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
end
end

# get_surface_params required by FluxCalculator (partitioned fluxes)
get_surface_params(sim::ClimaAtmosSimulation) = CAP.surface_fluxes_params(sim.integrator.p.params)

###
### ClimaAtmos.jl model-specific functions (not explicitly required by ClimaCoupler.jl)
###
"""
get_atmos_config(coupler_dict::Dict)
Returns the specified atmospheric configuration (`atmos_config_dict`) overwitten by arguments
in the coupler dictionary (`config_dict`).
"""
function get_atmos_config(coupler_dict)
atmos_config_file = coupler_dict["atmos_config_file"]
# override default or specified configs with coupler arguments, and set the correct atmos config_file
if isnothing(atmos_config_file)
@info "Using Atmos default configuration"
atmos_config = merge(CA.default_config_dict(), coupler_dict, Dict("config_file" => atmos_config_file))
else
@info "Using Atmos configuration from $atmos_config_file"
atmos_config = merge(
CA.override_default_config(joinpath(pkgdir(CA), atmos_config_file)),
coupler_dict,
Dict("config_file" => atmos_config_file),
)
end

# use coupler toml if atmos is not defined
atmos_toml_file = atmos_config["toml"]
coupler_toml_file = coupler_dict["coupler_toml_file"]
default_toml_file = "toml/default_coarse.toml"

toml_file = !isempty(atmos_toml_file) ? joinpath(pkgdir(CA), atmos_toml_file[1]) : nothing
toml_file = !isnothing(coupler_toml_file) ? joinpath(pkgdir(ClimaCoupler), coupler_toml_file) : toml_file
toml_file = isnothing(toml_file) ? joinpath(pkgdir(ClimaCoupler), default_toml_file) : toml_file

if !isnothing(toml_file)
@info "Overwriting Atmos parameters from $toml_file"
atmos_config = merge(atmos_config, Dict("toml" => [toml_file]))
end
return atmos_config
end


# flux calculation borrowed from atmos
"""
Expand Down Expand Up @@ -298,64 +359,6 @@ function calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::Cli
extrapolate_ρ_to_sfc.(Ref(thermo_params), ts_int, swap_space!(ones(axes(ts_int.ρ)), T_S))
end

"""
get_model_prog_state(sim::ClimaAtmosSimulation)
Extension of Checkpointer.get_model_prog_state to get the model state.
"""
function get_model_prog_state(sim::ClimaAtmosSimulation)
return sim.integrator.u
end

"""
get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
Extension of Interfacer.get_field to get the net TOA radiation, which is a sum of the
upward and downward longwave and shortwave radiation.
"""
function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
FT = eltype(atmos_sim.integrator.u)
# save radiation source
if atmos_sim.integrator.p.radiation.radiation_model != nothing
face_space = axes(atmos_sim.integrator.u.f)
nz_faces = length(ClimaCore.Spaces.vertical_topology(face_space).mesh.faces)

(; face_lw_flux_dn, face_lw_flux_up, face_sw_flux_dn, face_sw_flux_up) =
atmos_sim.integrator.p.radiation.radiation_model

LWd_TOA = ClimaCore.Fields.level(CA.RRTMGPI.array2field(FT.(face_lw_flux_dn), face_space), nz_faces - half)
LWu_TOA = ClimaCore.Fields.level(CA.RRTMGPI.array2field(FT.(face_lw_flux_up), face_space), nz_faces - half)
SWd_TOA = ClimaCore.Fields.level(CA.RRTMGPI.array2field(FT.(face_sw_flux_dn), face_space), nz_faces - half)
SWu_TOA = ClimaCore.Fields.level(CA.RRTMGPI.array2field(FT.(face_sw_flux_up), face_space), nz_faces - half)

return @. -(LWd_TOA + SWd_TOA - LWu_TOA - SWu_TOA) # [W/m^2]
else
return FT(0)
end
end

function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy})
thermo_params = get_thermo_params(atmos_sim)

ᶜS_ρq_tot = atmos_sim.integrator.p.precipitation.ᶜS_ρq_tot
ᶜts = atmos_sim.integrator.p.precomputed.ᶜts
ᶜΦ = atmos_sim.integrator.p.core.ᶜΦ

# return total energy and (if Microphysics0Moment) the energy lost due to precipitation removal
if atmos_sim.integrator.p.atmos.precip_model isa CA.Microphysics0Moment
ᶜS_ρq_tot = atmos_sim.integrator.p.precipitation.ᶜS_ρq_tot
ᶜts = atmos_sim.integrator.p.precomputed.ᶜts
ᶜΦ = atmos_sim.integrator.p.core.ᶜΦ
return atmos_sim.integrator.u.c.ρe_tot .-
ᶜS_ρq_tot .* CA.e_tot_0M_precipitation_sources_helper.(Ref(thermo_params), ᶜts, ᶜΦ) .*
atmos_sim.integrator.dt
else
return atmos_sim.integrator.u.c.ρe_tot
end
end

get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot

"""
dss_state!(sim::ClimaAtmosSimulation)
Expand Down
Loading

0 comments on commit 18856c1

Please sign in to comment.