From a643f22a9a62e548c5cae292b5d3aea546b97e08 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Tue, 23 Apr 2024 15:50:11 -0700 Subject: [PATCH] Preallocate thermal state --- docs/src/APIs/shared_utilities.md | 1 + .../Bucket/global_bucket_function.jl | 2 +- src/ClimaLand.jl | 14 +++++++++++ src/shared_utilities/drivers.jl | 24 +++++++++++++++---- src/shared_utilities/models.jl | 1 + src/standalone/Bucket/Bucket.jl | 1 + src/standalone/Vegetation/Canopy.jl | 9 ++++++- test/integrated/lsms.jl | 4 ++-- test/integrated/pond_soil_lsm.jl | 3 ++- test/standalone/Bucket/soil_bucket_tests.jl | 2 +- test/standalone/Soil/soil_test_3d.jl | 3 ++- test/standalone/Soil/soiltest.jl | 2 +- test/standalone/Vegetation/canopy_model.jl | 4 ++-- .../Vegetation/plant_hydraulics_test.jl | 5 ++-- 14 files changed, 59 insertions(+), 16 deletions(-) diff --git a/docs/src/APIs/shared_utilities.md b/docs/src/APIs/shared_utilities.md index eb748711b0..fe80503575 100644 --- a/docs/src/APIs/shared_utilities.md +++ b/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 diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl index b91dce5415..a06bc5107a 100644 --- a/experiments/standalone/Bucket/global_bucket_function.jl +++ b/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 diff --git a/src/ClimaLand.jl b/src/ClimaLand.jl index 623a457fca..8510775fb4 100644 --- a/src/ClimaLand.jl +++ b/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 diff --git a/src/shared_utilities/drivers.jl b/src/shared_utilities/drivers.jl index f10e38456f..b35e775f30 100644 --- a/src/shared_utilities/drivers.jl +++ b/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 diff --git a/src/shared_utilities/models.jl b/src/shared_utilities/models.jl index f9ccbcf9c6..b2a6e6b9aa 100644 --- a/src/shared_utilities/models.jl +++ b/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.", diff --git a/src/standalone/Bucket/Bucket.jl b/src/standalone/Bucket/Bucket.jl index bb629baba8..221923c461 100644 --- a/src/standalone/Bucket/Bucket.jl +++ b/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, diff --git a/src/standalone/Vegetation/Canopy.jl b/src/standalone/Vegetation/Canopy.jl index 24fc72ccf5..5a6fb20637 100644 --- a/src/standalone/Vegetation/Canopy.jl +++ b/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 diff --git a/test/integrated/lsms.jl b/test/integrated/lsms.jl index b9d66f2354..d104868a01 100644 --- a/test/integrated/lsms.jl +++ b/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) diff --git a/test/integrated/pond_soil_lsm.jl b/test/integrated/pond_soil_lsm.jl index 6eeb027647..a59407806d 100644 --- a/test/integrated/pond_soil_lsm.jl +++ b/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 diff --git a/test/standalone/Bucket/soil_bucket_tests.jl b/test/standalone/Bucket/soil_bucket_tests.jl index 56bab94896..d4c3310dcd 100644 --- a/test/standalone/Bucket/soil_bucket_tests.jl +++ b/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 diff --git a/test/standalone/Soil/soil_test_3d.jl b/test/standalone/Soil/soil_test_3d.jl index 1d991f82d7..e7009cbec8 100644 --- a/test/standalone/Soil/soil_test_3d.jl +++ b/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), diff --git a/test/standalone/Soil/soiltest.jl b/test/standalone/Soil/soiltest.jl index 5737916762..518ab1c123 100644 --- a/test/standalone/Soil/soiltest.jl +++ b/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 diff --git a/test/standalone/Vegetation/canopy_model.jl b/test/standalone/Vegetation/canopy_model.jl index f7291ef853..4135c233bb 100644 --- a/test/standalone/Vegetation/canopy_model.jl +++ b/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 diff --git a/test/standalone/Vegetation/plant_hydraulics_test.jl b/test/standalone/Vegetation/plant_hydraulics_test.jl index 7ec5da94f6..32e565f017 100644 --- a/test/standalone/Vegetation/plant_hydraulics_test.jl +++ b/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),