Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Generalize atmos_init to dry and no radiation cases #754

Merged
merged 1 commit into from
May 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,13 @@ steps:
agents:
slurm_mem: 20GB

- label: "Slabplanet: dry, no radiation"
key: "slabplanet_dry_norad"
command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/slabplanet_dry_norad.yml"
artifact_paths: "experiments/AMIP/output/slabplanet/slabplanet_dry_norad_artifacts/*"
agents:
slurm_mem: 20GB

- label: "Slabplanet: default with Float32"
key: "slabplanet_ft32"
command: "julia --color=yes --project=experiments/AMIP/ experiments/AMIP/coupler_driver.jl --config_file $CONFIG_PATH/slabplanet_ft32.yml"
Expand Down
18 changes: 18 additions & 0 deletions config/model_configs/slabplanet_dry_norad.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
anim: false
apply_limiter: false
conservation_softfail: true
dt: "200secs"
dt_cpl: 200
dt_save_to_sol: "9days"
energy_check: true
h_elem: 4
job_id: "slabplanet_dry_norad"
mode_name: "slabplanet"
moist: "dry"
mono_surface: true
precip_model: nothing
rad: nothing
run_name: "slabplanet_dry_norad"
t_end: "10days"
vert_diff: "true"
output_default_diagnostics: false
13 changes: 0 additions & 13 deletions experiments/AMIP/cli_options.jl
Original file line number Diff line number Diff line change
Expand Up @@ -185,19 +185,6 @@ function print_repl_script(str)
println(ib)
end

function time_to_seconds(s::String)
factor = Dict("secs" => 1, "mins" => 60, "hours" => 60 * 60, "days" => 60 * 60 * 24)
s == "Inf" && return Inf
if count(occursin.(keys(factor), Ref(s))) != 1
error("Bad format for flag $s. Examples: [`10secs`, `20mins`, `30hours`, `40days`]")
end
for match in keys(factor)
occursin(match, s) || continue
return parse(Float64, first(split(s, match))) * factor[match]
end
error("Uncaught case in computing time from given string.")
end

parsed_args_from_ARGS(ARGS, parsed_args = Dict()) = parsed_args_from_ARGS_string(strip(join(ARGS, " ")), parsed_args)

parsed_args_from_command_line_flags(str, parsed_args = Dict()) =
Expand Down
166 changes: 101 additions & 65 deletions experiments/AMIP/components/atmosphere/climaatmos.jl
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,14 @@ struct ClimaAtmosSimulation{P, Y, D, I} <: Interfacer.AtmosModelSimulation
end
Interfacer.name(::ClimaAtmosSimulation) = "ClimaAtmosSimulation"

function hasradiation(integrator)
return !isnothing(integrator.p.atmos.radiation_mode)
end

function hasmoisture(integrator)
return !(integrator.p.atmos.moisture_model isa CA.DryModel)
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 All @@ -34,27 +42,30 @@ function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
@warn("Running with ρe_int in coupled mode is not tested yet.", maxlog = 1)
end

# define shorter references for long variable Interfacer.names to increase readability
# define shorter references for long variable names to increase readability, and set to zero
ρ_flux_h_tot = integrator.p.precomputed.sfc_conditions.ρ_flux_h_tot
ρ_flux_q_tot = integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot
ᶠradiation_flux = integrator.p.radiation.ᶠradiation_flux
ρ_flux_uₕ = integrator.p.precomputed.sfc_conditions.ρ_flux_uₕ
col_integrated_rain = integrator.p.precipitation.col_integrated_rain
col_integrated_snow = integrator.p.precipitation.col_integrated_snow
ᶜS_ρq_tot = integrator.p.precipitation.ᶜS_ρq_tot
ᶜ3d_rain = integrator.p.precipitation.ᶜ3d_rain
ᶜ3d_snow = integrator.p.precipitation.ᶜ3d_snow

# set initial fluxes to zero
@. ρ_flux_h_tot = CC.Geometry.Covariant3Vector(FT(0.0))
@. ρ_flux_q_tot = CC.Geometry.Covariant3Vector(FT(0.0))
@. ᶠradiation_flux = CC.Geometry.WVector(FT(0))
ρ_flux_uₕ.components .= Ref(StaticArrays.SMatrix{1, 2}([FT(0), FT(0)]))
col_integrated_rain .= FT(0)
col_integrated_snow .= FT(0)
ᶜS_ρq_tot .= FT(0)
ᶜ3d_rain .= FT(0)
ᶜ3d_snow .= FT(0)

if hasmoisture(integrator)
ρ_flux_q_tot = integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot
col_integrated_rain = integrator.p.precipitation.col_integrated_rain
col_integrated_snow = integrator.p.precipitation.col_integrated_snow
ᶜS_ρq_tot = integrator.p.precipitation.ᶜS_ρq_tot
ᶜ3d_rain = integrator.p.precipitation.ᶜ3d_rain
ᶜ3d_snow = integrator.p.precipitation.ᶜ3d_snow
@. ρ_flux_q_tot = CC.Geometry.Covariant3Vector(FT(0.0))
col_integrated_rain .= FT(0)
col_integrated_snow .= FT(0)
ᶜS_ρq_tot .= FT(0)
ᶜ3d_rain .= FT(0)
ᶜ3d_snow .= FT(0)
end
if hasradiation(integrator)
ᶠradiation_flux = integrator.p.radiation.ᶠradiation_flux
@. ᶠradiation_flux = CC.Geometry.WVector(FT(0))
end

sim = ClimaAtmosSimulation(integrator.p.params, Y, spaces, integrator)

Expand Down Expand Up @@ -98,7 +109,7 @@ upward and downward longwave and shortwave radiation.
function Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
FT = eltype(atmos_sim.integrator.u)

if atmos_sim.integrator.p.radiation.radiation_model != nothing
if hasradiation(atmos_sim.integrator)
face_space = axes(atmos_sim.integrator.u.f)
nz_faces = length(CC.Spaces.vertical_topology(face_space).mesh.faces)

Expand All @@ -116,49 +127,74 @@ function Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_

return @. -(LWd_TOA + SWd_TOA - LWu_TOA - SWu_TOA)
else
return FT(0)
return [FT(0)]
end
end

function Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy})
thermo_params = get_thermo_params(atmos_sim)
integrator = atmos_sim.integrator
p = integrator.p

ᶜ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
if p.atmos.precip_model isa CA.Microphysics0Moment
ᶜts = p.precomputed.ᶜts
ᶜΦ = p.core.ᶜΦ
ᶜS_ρq_tot = p.precipitation.ᶜS_ρq_tot
thermo_params = get_thermo_params(atmos_sim)
return integrator.u.c.ρe_tot .-
ᶜS_ρq_tot .* CA.e_tot_0M_precipitation_sources_helper.(Ref(thermo_params), ᶜts, ᶜΦ) .* integrator.dt
else
return atmos_sim.integrator.u.c.ρe_tot
return integrator.u.c.ρe_tot
end
end

# helpers for get_field extensions, dipatchable on different moisture model options and radiation modes

col_integrated_rain(::CA.DryModel, integrator) = [eltype(integrator.u)(0)]
col_integrated_rain(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) =
integrator.p.precipitation.col_integrated_rain

col_integrated_snow(::CA.DryModel, integrator) = [eltype(integrator.u)(0)]
col_integrated_snow(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) =
integrator.p.precipitation.col_integrated_snow

surface_radiation_flux(::Nothing, integrator) = [eltype(integrator.u)(0)]
surface_radiation_flux(::CA.RRTMGPI.AbstractRRTMGPMode, integrator) =
CC.Fields.level(integrator.p.radiation.ᶠradiation_flux, CC.Utilities.half)

moisture_flux(::CA.DryModel, integrator) = [eltype(integrator.u)(0)]
moisture_flux(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) =
CC.Geometry.WVector.(integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot)

ρq_tot(::CA.DryModel, integrator) = [eltype(integrator.u)(0)]
ρq_tot(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) = integrator.u.c.ρq_tot

# extensions required by the Interfacer
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:air_density}) =
TD.air_density.(thermo_params, sim.integrator.p.precomputed.ᶜts)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:air_temperature}) =
TD.air_temperature.(thermo_params, sim.integrator.p.precomputed.ᶜts)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
sim.integrator.p.precipitation.col_integrated_rain
col_integrated_rain(sim.integrator.p.atmos.moisture_model, sim.integrator)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_sfc}) =
CC.Fields.level(sim.integrator.p.radiation.ᶠradiation_flux, CC.Utilities.half)
surface_radiation_flux(sim.integrator.p.atmos.radiation_mode, sim.integrator)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) =
sim.integrator.p.precipitation.col_integrated_snow
col_integrated_snow(sim.integrator.p.atmos.moisture_model, sim.integrator)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
CC.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_h_tot)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
CC.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot)
moisture_flux(sim.integrator.p.atmos.moisture_model, sim.integrator)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:thermo_state_int}) =
CC.Spaces.level(sim.integrator.p.precomputed.ᶜts, 1)
Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot

Interfacer.get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) =
ρq_tot(atmos_sim.integrator.p.atmos.moisture_model, atmos_sim.integrator)
function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, csf)
# note that this field is also being updated internally by the surface thermo state in ClimaAtmos
# if turbulent fluxes are calculated, to ensure consistency. In case the turbulent fluxes are not
# calculated, we update the field here.
sim.integrator.p.radiation.radiation_model.surface_temperature .= CA.RRTMGPI.field2array(csf.T_S)
end
# extensions required by FluxCalculator (partitioned fluxes)
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
CC.Spaces.level(CC.Fields.coordinate_field(sim.integrator.u.c).z, 1)
Expand All @@ -178,10 +214,6 @@ function Interfacer.update_field!(atmos_sim::ClimaAtmosSimulation, ::Val{:co2},
end
end
# extensions required by the Interfacer
function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, csf)
sim.integrator.p.radiation.radiation_model.surface_temperature .= CA.RRTMGPI.field2array(csf.T_S)
end

function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_direct_albedo}, field)
sim.integrator.p.radiation.radiation_model.direct_sw_surface_albedo .=
reshape(CA.RRTMGPI.field2array(field), 1, length(parent(field)))
Expand All @@ -196,7 +228,7 @@ function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fl
(; F_turb_energy, F_turb_moisture, F_turb_ρτxz, F_turb_ρτyz) = fields

Y = sim.integrator.u
surface_local_geometry = CC.Fields.level(CC.Fields.local_geometry_field(Y.f), CC.Fields.half)
surface_local_geometry = CC.Fields.level(CC.Fields.local_geometry_field(Y.f), CC.Utilities.half)
surface_normal = @. CA.C3(CA.unit_basis_vector_data(CA.C3, surface_local_geometry))

# get template objects for the contravariant components of the momentum fluxes (required by Atmos boundary conditions)
Expand Down Expand Up @@ -229,11 +261,16 @@ function FieldExchanger.update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbul
p = atmos_sim.integrator.p
t = atmos_sim.integrator.t

!p.radiation.idealized_insolation && CA.set_insolation_variables!(u, p, t)
!isempty(atmos_sim.integrator.p.radiation) &&
!p.radiation.idealized_insolation &&
CA.set_insolation_variables!(u, p, t)

if hasradiation(atmos_sim.integrator)
Interfacer.update_field!(atmos_sim, Val(:surface_direct_albedo), csf.surface_direct_albedo)
Interfacer.update_field!(atmos_sim, Val(:surface_diffuse_albedo), csf.surface_diffuse_albedo)
end

Interfacer.update_field!(atmos_sim, Val(:surface_direct_albedo), csf.surface_direct_albedo)
Interfacer.update_field!(atmos_sim, Val(:surface_diffuse_albedo), csf.surface_diffuse_albedo)
Interfacer.update_field!(atmos_sim, Val(:surface_temperature), csf)
!isempty(atmos_sim.integrator.p.radiation) && Interfacer.update_field!(atmos_sim, Val(:surface_temperature), csf)

if turbulent_fluxes isa FluxCalculator.PartitionedStateFluxes
Interfacer.update_field!(atmos_sim, Val(:turbulent_fluxes), csf)
Expand All @@ -258,12 +295,12 @@ FluxCalculator.get_surface_params(sim::ClimaAtmosSimulation) = CAP.surface_fluxe
### ClimaAtmos.jl model-specific functions (not explicitly required by ClimaCoupler.jl)
###
"""
get_atmos_config(coupler_dict::Dict)
get_atmos_config_dict(coupler_dict::Dict)

Returns the specified atmospheric configuration (`atmos_config_dict`) overwitten by arguments
in the coupler dictionary (`config_dict`).
Returns the specified atmospheric configuration (`atmos_config`) overwitten by arguments
in the coupler dictionary (`config_dict`). The returned dictionary will then be passed to CA.AtmosConfig().
"""
function get_atmos_config(coupler_dict)
function get_atmos_config_dict(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)
Expand Down Expand Up @@ -318,8 +355,12 @@ Sets up `surface_setup` as a `CC.Fields.Field` of `SurfaceState`s.
"""
function coupler_surface_setup(
::CoupledMoninObukhov,
p;
csf_sfc = (; T = nothing, z0m = nothing, z0b = nothing, beta = nothing, q_vap = nothing),
p,
T = nothing,
z0m = nothing,
z0b = nothing,
beta = nothing,
q_vap = nothing,
)

surface_state(z0m, z0b, T, beta, q_vap) = CA.SurfaceConditions.SurfaceState(;
Expand All @@ -328,7 +369,7 @@ function coupler_surface_setup(
beta,
q_vap,
)
surface_state_field = @. surface_state(csf_sfc.z0m, csf_sfc.z0b, csf_sfc.T, csf_sfc.beta, csf_sfc.q_vap)
surface_state_field = @. surface_state(z0m, z0b, T, beta, q_vap)
return surface_state_field
end

Expand All @@ -338,23 +379,18 @@ end
Returns a new `p` with the updated surface conditions.
"""
function get_new_cache(atmos_sim::ClimaAtmosSimulation, csf)
csf_sfc = (; T = csf.T_S, z0m = csf.z0m_S, z0b = csf.z0b_S, beta = csf.beta, q_vap = csf.q_sfc)
modified_atmos_cache(atmos_sim, csf_sfc)
end

"""
modified_atmos_cache(atmos_sim, csf_sfc)

Returns a new `p` with the updated surface conditions.
"""
function modified_atmos_cache(atmos_sim, csf_sfc)
if hasmoisture(atmos_sim.integrator)
csf_sfc = (csf.T_S, csf.z0m_S, csf.z0b_S, csf.beta, csf.q_sfc)
else
csf_sfc = (csf.T_S, csf.z0m_S, csf.z0b_S, csf.beta)
end

p = atmos_sim.integrator.p

coupler_sfc_setup = coupler_surface_setup(CoupledMoninObukhov(), p; csf_sfc = csf_sfc)
coupler_sfc_setup = coupler_surface_setup(CoupledMoninObukhov(), p, csf_sfc...)

p_names = propertynames(p)
p_values = map(x -> x == :sfc_setup ? coupler_sfc_setup : getproperty(p, x), p_names) # TODO: use merge here
p_values = map(x -> x == :sfc_setup ? coupler_sfc_setup : getproperty(p, x), p_names)

(; zip(p_names, p_values)...)
end
Expand Down
Loading
Loading