restructure component models

CliMA · Mar 8, 2024 · 1e66ec0 · 1e66ec0
1 parent 446f3cb
commit 1e66ec0
Show file tree

Hide file tree

Showing 13 changed files with 681 additions and 642 deletions.
diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md
@@ -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.
@@ -72,7 +72,7 @@ for the following properties:
 
 | Coupler name      | Description | Units |
 |-------------------|-------------|-------|
-| `air_density`       | air density at the surface of the atmosphere | kg m^-3 |
+| `air_density`       | air density of the atmosphere | kg m^-3 |
 | `air_temperature`   | air temperature at the surface of the atmosphere | K |
 | `energy`            | globally integrated energy | J |
 | `height_int`        | height at the first internal model level | m |
@@ -105,6 +105,9 @@ following properties:
 | `surface_temperature` | temperature over the combined surface space | K |
 | `turbulent_fluxes` | turbulent fluxes (note: only required when using `PartitionedStateFluxes` option - see our `FluxCalculator` module docs for more information) | W m^-2 |
 
+- `calculate_surface_air_density(atmos_sim::Interfacer.AtmosModelSimulation, T_S::Fields.Field)`:
+A function to return the air density of the atmosphere simulation
+extrapolated to the surface, with units of [kg m^-3].
 
 ### SurfaceModelSimulation - required functions
 Analogously to the `AtmosModelSimulation`, a `SurfaceModelSimulation`

diff --git a/.../components/atmosphere/climaatmos_init.jl → .../AMIP/components/atmosphere/climaatmos.jl b/.../components/atmosphere/climaatmos_init.jl → .../AMIP/components/atmosphere/climaatmos.jl
@@ -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
@@ -31,52 +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
-
-    # specify atmos output directory to be inside the coupler output directory
-    atmos_output_dir = joinpath(
-        coupler_dict["coupler_output_dir"],
-        joinpath(coupler_dict["mode_name"], coupler_dict["run_name"]),
-        "clima_atmos",
-    )
-    atmos_config = merge(atmos_config, Dict("output_dir" => atmos_output_dir))
-
-    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)
@@ -119,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)
+
+    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)
+    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
 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{:snow_precipitation}) = sim.integrator.p.precipitation.col_integrated_snow
 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}) =
@@ -140,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
@@ -167,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
 
@@ -197,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)
@@ -209,6 +218,69 @@ function update_sim!(atmos_sim::ClimaAtmosSimulation, csf, turbulent_fluxes)
     end
 end
 
+"""
+    calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::ClimaCore.Fields.Field)
+
+Extension for this function to calculate surface density.
+"""
+function calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::ClimaCore.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))
+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
+
+    # specify atmos output directory to be inside the coupler output directory
+    atmos_output_dir = joinpath(
+        coupler_dict["coupler_output_dir"],
+        joinpath(coupler_dict["mode_name"], coupler_dict["run_name"]),
+        "clima_atmos",
+    )
+    atmos_config = merge(atmos_config, Dict("output_dir" => atmos_output_dir))
+
+    return atmos_config
+end
+
 
 # flux calculation borrowed from atmos
 """
@@ -296,75 +368,6 @@ Returns the thermodynamic parameters from the atmospheric model simulation objec
 """
 get_thermo_params(sim::ClimaAtmosSimulation) = CAP.thermodynamics_params(sim.integrator.p.params)
 
-"""
-    calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::ClimaCore.Fields.Field)
-
-Extension for this  to to calculate surface density.
-"""
-function calculate_surface_air_density(atmos_sim::ClimaAtmosSimulation, T_S::ClimaCore.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))
-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)