Preallocate thermal state

.buildkite/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.3"
 manifest_format = "2.0"
-project_hash = "2bbfa2f663b4af45716f0653dbf2347a9562e36c"
+project_hash = "dae5b8a0ef6cad3a789e38991e9bafc7e71f7ab5"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "daf26bbdec60d9ca1c0003b70f389d821ddb4224"
@@ -373,7 +373,7 @@ weakdeps = ["CUDA", "Krylov"]
     KrylovExt = "Krylov"
 
 [[deps.ClimaLand]]
-deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities"]
+deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LazyArtifacts", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities"]
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.12.1"

.buildkite/Project.toml

@@ -20,6 +20,7 @@ HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"

docs/src/APIs/shared_utilities.md

@@ -77,7 +77,7 @@ ClimaLand.AbstractRadiativeDrivers
 ClimaLand.turbulent_fluxes
 ClimaLand.turbulent_fluxes_at_a_point
 ClimaLand.radiative_fluxes_at_a_point
-ClimaLand.construct_atmos_ts
+ClimaLand.set_atmos_ts!
 ClimaLand.surface_air_density
 ClimaLand.liquid_precipitation
 ClimaLand.snow_precipitation

docs/tutorials/standalone/Bucket/bucket_tutorial.jl

@@ -249,6 +249,7 @@ bucket_atmos = PrescribedAtmosphere(
     TimeVaryingInput(P_atmos),
     ref_time,
     h_atmos,
+    earth_param_set,
 );
 
 # Prescribed radiation -- a prescribed downwelling SW diurnal cycle, with a

docs/tutorials/standalone/Soil/evaporation.jl

@@ -76,7 +76,8 @@ atmos = PrescribedAtmosphere(
     TimeVaryingInput(q_atmos),
     TimeVaryingInput(P_atmos),
     ref_time,
-    h_atmos;
+    h_atmos,
+    earth_param_set;
     gustiness = gustiness,
 );
 

docs/tutorials/standalone/Soil/evaporation_gilat_loess.jl

@@ -106,7 +106,8 @@ atmos = PrescribedAtmosphere(
     TimeVaryingInput(q_atmos),
     TimeVaryingInput(P_atmos),
     ref_time,
-    h_atmos;
+    h_atmos,
+    earth_param_set;
     gustiness = gustiness,
 )
 

docs/tutorials/standalone/Soil/sublimation.jl

@@ -62,7 +62,8 @@ atmos = PrescribedAtmosphere(
     TimeVaryingInput(q_atmos),
     TimeVaryingInput(P_atmos),
     ref_time,
-    h_atmos;
+    h_atmos,
+    earth_param_set;
     gustiness = gustiness,
 );
 

experiments/integrated/fluxnet/met_drivers_FLUXNET.jl

@@ -145,7 +145,8 @@ atmos = ClimaLand.PrescribedAtmosphere(
     atmos_q,
     atmos_p,
     UTC_DATETIME[1],
-    atmos_h;
+    atmos_h,
+    earth_param_set;
     c_co2 = atmos_co2,
 )
 

experiments/integrated/performance/profile_allocations.jl

@@ -129,7 +129,8 @@ atmos = ClimaLand.PrescribedAtmosphere(
     atmos_q,
     atmos_p,
     ref_time,
-    atmos_h;
+    atmos_h,
+    earth_param_set;
 )
 function zenith_angle(
     t,

experiments/standalone/Biogeochemistry/experiment.jl

@@ -99,7 +99,8 @@ for (FT, tf) in ((Float32, 2 * dt), (Float64, tf))
         TimeVaryingInput(atmos_q),
         TimeVaryingInput(atmos_p),
         UTC_DATETIME,
-        atmos_h;
+        atmos_h,
+        earth_param_set;
         c_co2 = TimeVaryingInput(atmos_co2),
     )
 

experiments/standalone/Bucket/global_bucket_function.jl

@@ -111,6 +111,7 @@ bucket_atmos = PrescribedAtmosphere(
     TimeVaryingInput(P_atmos),
     ref_time,
     h_atmos,
+    earth_param_set,
 );
 
 # Prescribed radiation -- a prescribed downwelling SW diurnal cycle, with a

experiments/standalone/Bucket/global_bucket_staticmap.jl

@@ -182,6 +182,7 @@ bucket_atmos = PrescribedAtmosphere(
     P_atmos,
     ref_time,
     h_atmos,
+    earth_param_set,
 );
 
 # Prescribed radiation -- a prescribed downwelling SW diurnal cycle, with a

experiments/standalone/Bucket/global_bucket_temporalmap.jl

@@ -134,6 +134,7 @@ function setup_prob(t0, tf, Δt)
         TimeVaryingInput(P_atmos),
         ref_time,
         h_atmos,
+        earth_param_set,
     )
 
     # Prescribed radiation -- a prescribed downwelling SW diurnal cycle, with a

experiments/standalone/Snow/process_cdp.jl

@@ -2,6 +2,7 @@ using NCDatasets
 using ArtifactWrappers
 using Dates
 using ClimaLand: PrescribedAtmosphere, PrescribedRadiativeFluxes
+import ClimaLand.Parameters as LP
 using Thermodynamics
 using Statistics
 using Insolation
@@ -81,6 +82,7 @@ radiation = ClimaLand.PrescribedRadiativeFluxes(
 
 
 "Atmos"
+earth_param_set = LP.LandParameters(FT)
 liquid_precip = TimeVaryingInput(seconds, rainf; context)
 snow_precip = TimeVaryingInput(seconds, snowf; context)
 T_atmos = TimeVaryingInput(seconds, Tair; context)
@@ -98,6 +100,7 @@ atmos = PrescribedAtmosphere(
     P_atmos,
     ref_time,
     h_atmos,
+    earth_param_set,
 )
 # Meteorological data
 snow_af = ArtifactFile(

experiments/standalone/Soil/evaporation.jl

@@ -87,7 +87,8 @@ for (FT, tf) in ((Float32, 2 * dt), (Float64, tf))
         TimeVaryingInput(q_atmos),
         TimeVaryingInput(P_atmos),
         ref_time,
-        h_atmos;
+        h_atmos,
+        earth_param_set;
         gustiness = gustiness,
     )
     top_bc = ClimaLand.Soil.AtmosDrivenFluxBC(atmos, radiation)
@@ -210,7 +211,13 @@ for (FT, tf) in ((Float32, 2 * dt), (Float64, tf))
             )
             layer_thickness = Soil.dry_soil_layer_thickness.(S_l_sfc, S_c, d_ds)
             T_sfc = T_soil[i]
-            ts_in = ClimaLand.construct_atmos_ts(top_bc.atmos, p, thermo_params)
+            ts_in =
+                Thermodynamics.PhaseEquil_pTq.(
+                    thermo_params,
+                    p.drivers.P,
+                    p.drivers.T,
+                    p.drivers.q,
+                )
             ρ_sfc = compute_ρ_sfc.(thermo_params, ts_in, T_sfc)
             q_sat =
                 Thermodynamics.q_vap_saturation_generic.(

lib/ClimaLandSimulations/src/Fluxnet/fluxnet_utilities/make_drivers.jl

@@ -133,7 +133,8 @@ function make_drivers(site_ID, setup, config, params, context)
         atmos_q,
         atmos_p,
         UTC_DATETIME[1],
-        config.atmos_h;
+        config.atmos_h,
+        earth_param_set;
         c_co2 = atmos_co2,
     )
 

src/shared_utilities/drivers.jl

@@ -15,7 +15,7 @@ export AbstractAtmosphericDrivers,
     PrescribedRadiativeFluxes,
     CoupledRadiativeFluxes,
     compute_ρ_sfc,
-    construct_atmos_ts,
+    set_atmos_ts!,
     turbulent_fluxes,
     net_radiation,
     turbulent_fluxes_at_a_point,
@@ -63,6 +63,7 @@ struct PrescribedAtmosphere{
     RA <: Union{Nothing, AbstractTimeVaryingInput},
     CA <: Union{Nothing, AbstractTimeVaryingInput},
     DT,
+    TP,
 } <: AbstractAtmosphericDrivers{FT}
     "Precipitation (m/s) function of time: positive by definition"
     liquid_precip::LP
@@ -84,6 +85,8 @@ struct PrescribedAtmosphere{
     h::FT
     "Minimum wind speed (gustiness; m/s)"
     gustiness::FT
+    "Thermodynamic parameters"
+    thermo_params::TP
     function PrescribedAtmosphere(
         liquid_precip,
         snow_precip,
@@ -92,12 +95,19 @@ struct PrescribedAtmosphere{
         q,
         P,
         ref_time,
-        h::FT;
+        h::FT,
+        earth_param_set;
         gustiness = FT(1),
         c_co2 = TimeVaryingInput((t) -> 4.2e-4),
     ) where {FT}
+        thermo_params = LP.thermodynamic_parameters(earth_param_set)
         args = (liquid_precip, snow_precip, T, u, q, P, c_co2, ref_time)
-        return new{typeof(h), typeof.(args)...}(args..., h, gustiness)
+        return new{typeof(h), typeof.(args)..., typeof(thermo_params)}(
+            args...,
+            h,
+            gustiness,
+            thermo_params,
+        )
     end
 end
 
@@ -152,30 +162,19 @@ function compute_ρ_sfc(thermo_params, ts_in, T_sfc)
 end
 
 """
-    construct_atmos_ts(
-        atmos::PrescribedAtmosphere,
-        p,
-        thermo_params,
-    )
+    set_atmos_ts!(ts_in, atmos::PrescribedAtmosphere{FT}, p)
 
-A helper function which constructs a Clima.Thermodynamics
-thermodynamic state given a PrescribedAtmosphere, the cache `p`,
- and a set of Clima.Thermodynamics
-parameters thermo_params.
+Fill the pre-allocated ts_in `Field` with a thermodynamic state computed from the
+atmosphere.
 """
-function construct_atmos_ts(
-    atmos::PrescribedAtmosphere{FT},
-    p,
-    thermo_params,
-) where {FT}
+function set_atmos_ts!(ts_in, atmos::PrescribedAtmosphere{FT}, p) where {FT}
     P = p.drivers.P
     T = p.drivers.T
     q = p.drivers.q
-    ts_in = Thermodynamics.PhaseEquil_pTq.(thermo_params, P, T, q)
-    return ts_in
+    ts_in .= Thermodynamics.PhaseEquil_pTq.(atmos.thermo_params, P, T, q)
+    return nothing
 end
 
-
 """
     turbulent_fluxes(atmos::PrescribedAtmosphere,
                    model::AbstractModel,
@@ -206,9 +205,6 @@ function turbulent_fluxes(
     h_sfc = surface_height(model, Y, p)
     r_sfc = surface_resistance(model, Y, p, t)
     d_sfc = displacement_height(model, Y, p)
-    thermo_params =
-        LP.thermodynamic_parameters(model.parameters.earth_param_set)
-    ts_air = construct_atmos_ts(atmos, p, thermo_params)
     u_air = p.drivers.u
     h_air = atmos.h
 
@@ -220,7 +216,7 @@ function turbulent_fluxes(
         h_sfc,
         r_sfc,
         d_sfc,
-        ts_air,
+        p.drivers.thermal_state,
         u_air,
         h_air,
         atmos.gustiness,
@@ -518,8 +514,7 @@ function surface_air_density(
 )
     thermo_params =
         LP.thermodynamic_parameters(model.parameters.earth_param_set)
-    ts_in = construct_atmos_ts(atmos, p, thermo_params)
-    return compute_ρ_sfc.(thermo_params, ts_in, T_sfc)
+    return compute_ρ_sfc.(thermo_params, p.drivers.thermal_state, T_sfc)
 end
 
 
@@ -662,8 +657,9 @@ horizontal wind speed `u`, specific humidity `q`, and CO2 concentration
 `c_co2`.
 """
 function initialize_drivers(a::PrescribedAtmosphere{FT}, coords) where {FT}
-    keys = (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2)
-    types = ([FT for k in keys]...,)
+    keys = (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2, :thermal_state)
+    # The thermal state is a different type
+    types = ([FT for k in keys[1:(end - 1)]]..., Thermodynamics.PhaseEquil{FT})
     domain_names = ([:surface for k in keys]...,)
     model_name = :drivers
     # intialize_vars packages the variables as a named tuple,
@@ -813,6 +809,7 @@ function make_update_drivers(a::PrescribedAtmosphere{FT}) where {FT}
         evaluate!(p.drivers.u, a.u, t)
         evaluate!(p.drivers.q, a.q, t)
         evaluate!(p.drivers.c_co2, a.c_co2, t)
+        set_atmos_ts!(p.drivers.thermal_state, a, p)
     end
     return update_drivers!
 end

src/shared_utilities/models.jl

@@ -191,7 +191,6 @@ driver variables in the cache.
 """
 add_drivers_to_cache(p, model::AbstractModel) = p
 
-
 """
     make_imp_tendency(model::AbstractImExModel)
 
@@ -387,6 +386,16 @@ are required, `p` is returned unchanged.
 """
 function add_drivers_to_cache(p::NamedTuple, model::AbstractModel, coords)
     (atmos, radiation) = get_drivers(model)
+    if hasproperty(model, :parameters) &&
+       hasproperty(model.parameters, :earth_param_set) &&
+       !isnothing(atmos)
+        if LP.thermodynamic_parameters(model.parameters.earth_param_set) !=
+           atmos.thermo_params
+            error(
+                "earth_param_set is inconsistent between the model and the atmosphere",
+            )
+        end
+    end
     driver_nt = initialize_drivers(atmos, radiation, coords)
     if driver_nt == (;)
         return p

src/standalone/Bucket/Bucket.jl

@@ -26,7 +26,6 @@ using ClimaLand:
     snow_precipitation,
     turbulent_fluxes,
     net_radiation,
-    construct_atmos_ts,
     compute_ρ_sfc,
     AbstractExpModel,
     heaviside,

test/integrated/soil_energy_hydrology_biogeochemistry.jl

@@ -92,7 +92,8 @@ for FT in (Float32, Float64)
             TimeVaryingInput(atmos_q),
             TimeVaryingInput(atmos_p),
             UTC_DATETIME,
-            atmos_h;
+            atmos_h,
+            earth_param_set;
             c_co2 = TimeVaryingInput(atmos_co2),
         )
 
@@ -116,7 +117,7 @@ for FT in (Float32, Float64)
         )
         Y, p, coords = initialize(model)
         @test propertynames(p.drivers) ==
-              (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2)
+              (:P_liq, :P_snow, :T, :P, :u, :q, :c_co2, :thermal_state)
         function init_soil!(Y, z, params)
             ν = params.ν
             FT = eltype(Y.soil.ϑ_l)

test/shared_utilities/drivers.jl

@@ -1,5 +1,6 @@
 import ClimaComms
 @static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
+import ClimaLand.Parameters as LP
 using ClimaCore
 using Test
 using StaticArrays
@@ -84,8 +85,19 @@ end
 end
 
 @testset "Driver update functions" begin
+    earth_param_set = LP.LandParameters(FT)
     f = TimeVaryingInput((t) -> 10.0)
-    pa = ClimaLand.PrescribedAtmosphere(f, f, f, f, f, f, f, FT(1);)
+    pa = ClimaLand.PrescribedAtmosphere(
+        f,
+        f,
+        f,
+        f,
+        f,
+        f,
+        f,
+        FT(1),
+        earth_param_set,
+    )
     pr = ClimaLand.PrescribedRadiativeFluxes(FT, f, f, f)
     coords = (; surface = [1])
     p = (; drivers = ClimaLand.initialize_drivers(nothing, nothing, coords))