Preallocate thermal state

docs/src/APIs/shared_utilities.md

@@ -77,6 +77,7 @@ 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

experiments/standalone/Bucket/global_bucket_function.jl

@@ -161,7 +161,7 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-@time sol =
+ClimaComms.@time sol =
     SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
 
 # Interpolate to grid

src/ClimaLand.jl

@@ -106,10 +106,24 @@ function initialize_auxiliary(
     end
     for domain in unique(domains_list)
         p = add_dss_buffer_to_aux(p, domain)
+        p = add_scratch_to_aux(p, domain)
     end
     return p
 end
 
+function add_scratch_to_aux(p, domain)
+    hasproperty(domain.space, :surface) || error("Surface domain required")
+    FT = Spaces.undertype(domain.space.surface)
+    scratch = (;
+        # Thermal state on the surface
+        thermal_state = Fields.Field(
+            Thermodynamics.PhaseEquil{FT},
+            domain.space.surface,
+        )
+    )
+    return merge(p, (; scratch))
+end
+
 """
     initialize_lsm_aux(land::AbstractLandModel) end
 

src/shared_utilities/drivers.jl

@@ -16,6 +16,7 @@ export AbstractAtmosphericDrivers,
     CoupledRadiativeFluxes,
     compute_ρ_sfc,
     construct_atmos_ts,
+    set_atmos_ts!,
     turbulent_fluxes,
     net_radiation,
     turbulent_fluxes_at_a_point,
@@ -175,6 +176,21 @@ function construct_atmos_ts(
     return ts_in
 end
 
+"""
+Fill the pre-allocated ts_in `Field` with a thermodynamic state.
+"""
+function set_atmos_ts!(
+    ts_in,
+    atmos::PrescribedAtmosphere{FT},
+    p,
+    thermo_params,
+) 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 nothing
+end
 
 """
     turbulent_fluxes(atmos::PrescribedAtmosphere,
@@ -208,7 +224,7 @@ function turbulent_fluxes(
     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)
+    set_atmos_ts!(p.scratch.thermal_state, atmos, p, thermo_params)
     u_air = p.drivers.u
     h_air = atmos.h
 
@@ -220,7 +236,7 @@ function turbulent_fluxes(
         h_sfc,
         r_sfc,
         d_sfc,
-        ts_air,
+        p.scratch.thermal_state,
         u_air,
         h_air,
         atmos.gustiness,
@@ -518,8 +534,8 @@ 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)
+    set_atmos_ts!(p.scratch.thermal_state, atmos, p, thermo_params)
+    return compute_ρ_sfc.(thermo_params, p.scratch.thermal_state, T_sfc)
 end
 
 

src/shared_utilities/models.jl

@@ -355,6 +355,7 @@ function initialize_auxiliary(model::AbstractModel{FT}, state) where {FT}
     )
     if :domain ∈ propertynames(model)
         p = add_dss_buffer_to_aux(p, model.domain)
+        p = add_scratch_to_aux(p, model.domain)
     else
         error(
             "Your model does not contain a domain. If this is intended, you will need a new method of initialize_auxiliary.",

src/standalone/Bucket/Bucket.jl

@@ -27,6 +27,7 @@ using ClimaLand:
     turbulent_fluxes,
     net_radiation,
     construct_atmos_ts,
+    set_atmos_ts!,
     compute_ρ_sfc,
     AbstractExpModel,
     heaviside,

src/standalone/Vegetation/Canopy.jl

@@ -354,7 +354,14 @@ function initialize_auxiliary(model::CanopyModel{FT}, coords) where {FT}
     # `p_state_list` contains `nothing` for components with no auxiliary
     #  variables, which we need to filter out before constructing `p`
     p = (; name(model) => filter_nt(NamedTuple{components}(p_state_list)))
-    p = ClimaLand.add_dss_buffer_to_aux(p, model.domain)
+    if :domain ∈ propertynames(model)
+        p = ClimaLand.add_dss_buffer_to_aux(p, model.domain)
+        p = ClimaLand.add_scratch_to_aux(p, model.domain)
+    else
+        error(
+            "Your model does not contain a domain. If this is intended, you will need a new method of initialize_auxiliary.",
+        )
+    end
     return p
 end
 

test/integrated/lsms.jl

@@ -90,7 +90,7 @@ for FT in (Float32, Float64)
         m2 = DummyModel2{FT}(d)
         m = DummyModel{FT}(m1, m2)
         Y, p, cds = ClimaLand.initialize(m)
-        @test propertynames(p) == (:i1, :m1, :m2)
+        @test propertynames(p) == (:i1, :m1, :m2, :scratch)
         exp_tendency! = make_exp_tendency(m)
         dY = similar(Y)
         exp_tendency!(dY, Y, p, FT(0))
@@ -137,7 +137,7 @@ for FT in (Float32, Float64)
         m2 = DummyModel4{FT}(d)
         m = DummyModelB{FT}(m1, m2)
         Y, p, cds = ClimaLand.initialize(m)
-        @test propertynames(p) == (:m1, :m2)
+        @test propertynames(p) == (:m1, :m2, :scratch)
         function ClimaLand.make_update_aux(::DummyModel3{FT}) where {FT}
             function update_aux!(p, Y, t)
                 p.m1.b .= FT(10.0)

test/integrated/pond_soil_lsm.jl

@@ -86,6 +86,7 @@ for FT in (Float32, Float64)
                     :surface_water,
                     :dss_buffer_3d,
                     :dss_buffer_2d,
+                    :scratch,
                 )
                 @test typeof(p.dss_buffer_3d) == typeof(
                     ClimaCore.Spaces.create_dss_buffer(
@@ -144,7 +145,7 @@ for FT in (Float32, Float64)
                       ClimaLand.lsm_aux_domain_names(land)
             else
                 @test propertynames(p) ==
-                      (:soil_infiltration, :soil, :surface_water)
+                      (:soil_infiltration, :soil, :surface_water, :scratch)
             end
         end
     end

test/standalone/Bucket/soil_bucket_tests.jl

@@ -120,7 +120,7 @@ for FT in (Float32, Float64)
                     ),
                 )
             else
-                @test propertynames(p) == (:bucket, :drivers)
+                @test propertynames(p) == (:bucket, :scratch, :drivers)
             end
 
 

test/standalone/Soil/soil_test_3d.jl

@@ -71,7 +71,8 @@ for FT in (Float32, Float64)
 
         Y, p, coords = initialize(soil)
         # test that the dss buffer was added
-        @test propertynames(p) == (:soil, :dss_buffer_3d, :dss_buffer_2d)
+        @test propertynames(p) ==
+              (:soil, :dss_buffer_3d, :dss_buffer_2d, :scratch)
         @test typeof(p.dss_buffer_3d) == typeof(
             ClimaCore.Spaces.create_dss_buffer(
                 ClimaCore.Fields.zeros(soil_domain.space.subsurface),

test/standalone/Soil/soiltest.jl

@@ -46,7 +46,7 @@ for FT in (Float32, Float64)
         # Here we do not add a dss buffer: check that `initialize`
         # has not added a buffer to `p` and that it only contains the
         # `soil` variables.
-        @test propertynames(p) == (:soil,)
+        @test propertynames(p) == (:soil, :scratch)
         # specify ICs
         function init_soil!(Ysoil, z, params)
             (; ν, hydrology_cm, θ_r, S_s) = params

test/standalone/Vegetation/canopy_model.jl

@@ -191,7 +191,7 @@ for FT in (Float32, Float64)
         # Check that structure of Y is value (will error if not)
         @test !isnothing(zero(Y))
         @test typeof(canopy.energy) == PrescribedCanopyTempModel{FT}
-        @test propertynames(p) == (:canopy, :drivers)
+        @test propertynames(p) == (:canopy, :scratch, :drivers)
         for component in ClimaLand.Canopy.canopy_components(canopy)
             # Only hydraulics has a prognostic variable
             if component == :hydraulics
@@ -666,7 +666,7 @@ for FT in (Float32, Float64)
 
         # Check that structure of Y is value (will error if not)
         @test !isnothing(zero(Y))
-        @test propertynames(p) == (:canopy, :drivers)
+        @test propertynames(p) == (:canopy, :scratch, :drivers)
         for component in ClimaLand.Canopy.canopy_components(canopy)
             # Only hydraulics has a prognostic variable
             if component == :hydraulics

test/standalone/Vegetation/plant_hydraulics_test.jl

@@ -347,9 +347,10 @@ for FT in (Float32, Float64)
 
             Y, p, coords = initialize(model)
             if typeof(domain) <: ClimaLand.Domains.Point
-                @test propertynames(p) == (:canopy, :drivers)
+                @test propertynames(p) == (:canopy, :scratch, :drivers)
             elseif typeof(domain) <: ClimaLand.Domains.Plane
-                @test propertynames(p) == (:canopy, :dss_buffer_2d, :drivers)
+                @test propertynames(p) ==
+                      (:canopy, :dss_buffer_2d, :scratch, :drivers)
                 @test typeof(p.dss_buffer_2d) == typeof(
                     ClimaCore.Spaces.create_dss_buffer(
                         ClimaCore.Fields.zeros(domain.space.surface),