Skip to content

Commit

Permalink
exp/AMIP
Browse files Browse the repository at this point in the history
  • Loading branch information
@juliasloan25
juliasloan25 committed Apr 9, 2024
1 parent 11cc032 commit 6c81b77
Show file tree
Hide file tree
Showing 15 changed files with 378 additions and 375 deletions.
112 changes: 51 additions & 61 deletions experiments/AMIP/components/atmosphere/climaatmos.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,36 +2,23 @@
import StaticArrays
import Statistics
import LinearAlgebra

import ClimaAtmos as CA
import ClimaCore as CC
import SurfaceFluxes as SF

import ClimaCoupler.Interfacer: AtmosModelSimulation
import ClimaCoupler.FluxCalculator:
atmos_turbulent_fluxes!,
calculate_surface_air_density,
PartitionedStateFluxes,
extrapolate_ρ_to_sfc,
get_surface_params,
water_albedo_from_atmosphere!
import ClimaCoupler.Interfacer: get_field, update_field!, name, step!, reinit!
import ClimaCoupler.Checkpointer: get_model_prog_state
import ClimaCoupler.FieldExchanger: update_sim!
import ClimaCoupler.Utilities: swap_space!
import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities

include("climaatmos_extra_diags.jl")

###
### Functions required by ClimaCoupler.jl for an AtmosModelSimulation
###
struct ClimaAtmosSimulation{P, Y, D, I} <: AtmosModelSimulation
struct ClimaAtmosSimulation{P, Y, D, I} <: Interfacer.AtmosModelSimulation
params::P
Y_init::Y
domain::D
integrator::I
end
name(::ClimaAtmosSimulation) = "ClimaAtmosSimulation"
Interfacer.name(::ClimaAtmosSimulation) = "ClimaAtmosSimulation"

function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
# By passing `parsed_args` to `AtmosConfig`, `parsed_args` overwrites the default atmos config
Expand All @@ -47,7 +34,7 @@ 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 names to increase readability
# define shorter references for long variable Interfacer.names to increase readability
ρ_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
Expand Down Expand Up @@ -77,21 +64,21 @@ function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
end

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

"""
get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_toa})
Interfacer.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})
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
Expand All @@ -116,7 +103,7 @@ function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux
end
end

function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy})
function Interfacer.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
Expand All @@ -137,32 +124,35 @@ function get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy})
end

# extensions required by the Interfacer
get_field(sim::ClimaAtmosSimulation, ::Val{:air_density}) =
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}) =
Interfacer.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
get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_sfc}) =
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:liquid_precipitation}) =
sim.integrator.p.precipitation.col_integrated_rain
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:radiative_energy_flux_sfc}) =
CC.Fields.level(sim.integrator.p.radiation.ᶠradiation_flux, CC.Utilities.half)
get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) = sim.integrator.p.precipitation.col_integrated_snow
get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:snow_precipitation}) =
sim.integrator.p.precipitation.col_integrated_snow
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_energy_flux}) =
CC.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_h_tot)
get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:turbulent_moisture_flux}) =
CC.Geometry.WVector.(sim.integrator.p.precomputed.sfc_conditions.ρ_flux_q_tot)
get_field(sim::ClimaAtmosSimulation, ::Val{:thermo_state_int}) = CC.Spaces.level(sim.integrator.p.precomputed.ᶜts, 1)
get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot
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

# extensions required by FluxCalculator (partitioned fluxes)
get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:height_int}) =
CC.Spaces.level(CC.Fields.coordinate_field(sim.integrator.u.c).z, 1)
get_field(sim::ClimaAtmosSimulation, ::Val{:height_sfc}) =
Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:height_sfc}) =
CC.Spaces.level(CC.Fields.coordinate_field(sim.integrator.u.f).z, CC.Utilities.half)
function get_field(sim::ClimaAtmosSimulation, ::Val{:uv_int})
function Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:uv_int})
uₕ_int = CC.Geometry.UVVector.(CC.Spaces.level(sim.integrator.u.c.uₕ, 1))
return @. StaticArrays.SVector(uₕ_int.components.data.:1, uₕ_int.components.data.:2)
end

function update_field!(atmos_sim::ClimaAtmosSimulation, ::Val{:co2}, field)
function Interfacer.update_field!(atmos_sim::ClimaAtmosSimulation, ::Val{:co2}, field)
if atmos_sim.integrator.p.atmos.radiation_mode isa CA.RRTMGPI.GrayRadiation
@warn("Gray radiation model initialized, skipping CO2 update", maxlog = 1)
return
Expand All @@ -171,21 +161,21 @@ function update_field!(atmos_sim::ClimaAtmosSimulation, ::Val{:co2}, field)
end
end
# extensions required by the Interfacer
function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_temperature}, csf)
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 update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_direct_albedo}, field)
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)))
end

function update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_diffuse_albedo}, field)
function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:surface_diffuse_albedo}, field)
sim.integrator.p.radiation.radiation_model.diffuse_sw_surface_albedo .=
reshape(CA.RRTMGPI.field2array(field), 1, length(parent(field)))
end

function update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fields)
function Interfacer.update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fields)
(; F_turb_energy, F_turb_moisture, F_turb_ρτxz, F_turb_ρτyz) = fields

Y = sim.integrator.u
Expand All @@ -199,8 +189,8 @@ function update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fiel
sim.integrator.p.precomputed.sfc_conditions.ρ_flux_uₕ .= (
surface_normal.CA .⊗
CA.C12.(
swap_space!(ones(axes(vec_ct12_ct1)), F_turb_ρτxz) .* vec_ct12_ct1 .+
swap_space!(ones(axes(vec_ct12_ct2)), F_turb_ρτyz) .* vec_ct12_ct2,
Utilities.swap_space!(ones(axes(vec_ct12_ct1)), F_turb_ρτxz) .* vec_ct12_ct1 .+
Utilities.swap_space!(ones(axes(vec_ct12_ct2)), F_turb_ρτyz) .* vec_ct12_ct2,
surface_local_geometry,
)
)
Expand All @@ -213,32 +203,32 @@ function update_field!(sim::ClimaAtmosSimulation, ::Val{:turbulent_fluxes}, fiel
end

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

function update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
update_field!(atmos_sim, Val(:surface_direct_albedo), csf.surface_direct_albedo)
update_field!(atmos_sim, Val(:surface_diffuse_albedo), csf.surface_diffuse_albedo)
update_field!(atmos_sim, Val(:surface_temperature), csf)
function FieldExchanger.update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
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)

if turbulent_fluxes isa PartitionedStateFluxes
update_field!(atmos_sim, Val(:turbulent_fluxes), csf)
if turbulent_fluxes isa FluxCalculator.PartitionedStateFluxes
Interfacer.update_field!(atmos_sim, Val(:turbulent_fluxes), csf)
end
end

"""
calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::CC.Fields.Field)
FluxCalculator.calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::CC.Fields.Field)
Extension for this function to calculate surface density.
"""
function calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::CC.Fields.Field)
function FluxCalculator.calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::CC.Fields.Field)
thermo_params = get_thermo_params(atmos_sim)
ts_int = get_field(atmos_sim, Val(:thermo_state_int))
extrapolate_ρ_to_sfc.(Ref(thermo_params), ts_int, swap_space!(ones(axes(ts_int.ρ)), T_S))
ts_int = Interfacer.get_field(atmos_sim, Val(:thermo_state_int))
FluxCalculator.extrapolate_ρ_to_sfc.(Ref(thermo_params), ts_int, Utilities.swap_space!(ones(axes(ts_int.ρ)), T_S))
end

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

###
### ClimaAtmos.jl model-specific functions (not explicitly required by ClimaCoupler.jl)
Expand Down Expand Up @@ -346,12 +336,12 @@ function modified_atmos_cache(atmos_sim, csf_sfc)
end

"""
atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
FluxCalculator.atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
Computes turbulent surface fluxes using ClimaAtmos's `update_surface_conditions!`. This
requires that we define a new temporary parameter Tuple, `new_p`, and save the new surface state
in it. We do not want `new_p` to live in the atmospheric model permanently, because that would also
trigger flux calculation during Atmos `step!`. We only want to trigger this once per coupling
trigger flux calculation during Atmos `Interfacer.step!`. We only want to trigger this once per coupling
timestep from ClimaCoupler.
For debigging atmos, we can set the following atmos defaults:
Expand All @@ -360,7 +350,7 @@ For debigging atmos, we can set the following atmos defaults:
csf.beta .= 1
csf = merge(csf, (;q_sfc = nothing))
"""
function atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)
function FluxCalculator.atmos_turbulent_fluxes!(atmos_sim::ClimaAtmosSimulation, csf)

if isnothing(atmos_sim.integrator.p.sfc_setup) # trigger flux calculation if not done in Atmos internally
new_p = get_new_cache(atmos_sim, csf)
Expand Down Expand Up @@ -394,11 +384,11 @@ function dss_state!(sim::ClimaAtmosSimulation)
end

"""
water_albedo_from_atmosphere!(atmos_sim::ClimaAtmosSimulation, direct_albedo::CC.Fields.Field, diffuse_albedo::CC.Fields.Field)
FluxCalculator.water_albedo_from_atmosphere!(atmos_sim::ClimaAtmosSimulation, direct_albedo::CC.Fields.Field, diffuse_albedo::CC.Fields.Field)
Extension to calculate the water surface albedo from wind speed and insolation. It can be used for prescribed ocean and lakes.
"""
function water_albedo_from_atmosphere!(
function FluxCalculator.water_albedo_from_atmosphere!(
atmos_sim::ClimaAtmosSimulation,
direct_albedo::CC.Fields.Field,
diffuse_albedo::CC.Fields.Field,
Expand Down
1 change: 0 additions & 1 deletion experiments/AMIP/components/atmosphere/climaatmos_extra_diags.jl
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
# these extensions add extra diagnostics to the atmos model output
import ClimaAtmos.Diagnostics as CAD
import ClimaAtmos.Diagnostics: add_diagnostic_variable!

"""
add_diagnostic_variable!(short_name::String, long_name::String, standard_name::String, units::String, comments::String, compute!::Function)
Expand Down
Loading

0 comments on commit 6c81b77

Please sign in to comment.