From 43801de84fd705cea3286cb2979dbe83f110e402 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Tue, 26 Nov 2024 12:07:41 -0800 Subject: [PATCH] add wrapper struct for type inference This is required as a workaround to the type inference failure detailed in https://github.com/CliMA/ClimaCore.jl/issues/2065. The failure appears when a function is broadcast over one field and two structs; the types are too complex and inference fails. This workaround introduces a wrapper struct to contain the two struct inputs, so that the function is broadcast over one field and only one struct, and the types are easier to infer. --- NEWS.md | 4 +- src/standalone/Soil/energy_hydrology.jl | 12 +-- .../Soil/soil_heat_parameterizations.jl | 40 ++++++++-- .../standalone/Soil/soil_parameterizations.jl | 78 ++++++++++++++++--- 4 files changed, 109 insertions(+), 25 deletions(-) diff --git a/NEWS.md b/NEWS.md index d651efec36..5d4bf6f104 100644 --- a/NEWS.md +++ b/NEWS.md @@ -3,7 +3,9 @@ ClimaLand.jl Release Notes main -------- - +- Run unit tests on GPU, and update code for GPU compatibility + (including workaround for ClimaCore type inference failure) + PR[#739](https://github.com/CliMA/ClimaLand.jl/pull/739) v0.15.6 ------- diff --git a/src/standalone/Soil/energy_hydrology.jl b/src/standalone/Soil/energy_hydrology.jl index 929ddb3b57..04a995ece9 100644 --- a/src/standalone/Soil/energy_hydrology.jl +++ b/src/standalone/Soil/energy_hydrology.jl @@ -623,6 +623,11 @@ function ClimaLand.source!( _ρ_l = FT(LP.ρ_cloud_liq(earth_param_set)) _ρ_i = FT(LP.ρ_cloud_ice(earth_param_set)) Δz = model.domain.fields.Δz # center face distance + + # Wrap hydrology and earth parameters in one struct to avoid type inference failure + hydrology_earth_params = + ClimaLand.Soil.HydrologyEarthParameters.(hydrology_cm, earth_param_set) + @. dY.soil.ϑ_l += -phase_change_source( p.soil.θ_l, @@ -640,8 +645,7 @@ function ClimaLand.source!( ), ν, θ_r, - hydrology_cm, - earth_param_set, + hydrology_earth_params, ) @. dY.soil.θ_i += (_ρ_l / _ρ_i) * phase_change_source( @@ -660,12 +664,10 @@ function ClimaLand.source!( ), ν, θ_r, - hydrology_cm, - earth_param_set, + hydrology_earth_params, ) end - """ SoilSublimation{FT} <: AbstractSoilSource{FT} diff --git a/src/standalone/Soil/soil_heat_parameterizations.jl b/src/standalone/Soil/soil_heat_parameterizations.jl index dc6fb35427..bcccdb50e9 100644 --- a/src/standalone/Soil/soil_heat_parameterizations.jl +++ b/src/standalone/Soil/soil_heat_parameterizations.jl @@ -31,12 +31,11 @@ end τ::FT, ν::FT, θ_r::FT, - hydrology_cm::C, - earth_param_set::EP, - ) where {FT, EP, C} + hydrology_earth_params::HEP + ) where {FT, HEP} Returns the source term (1/s) used for converting liquid water and ice into each other during phase changes. Note that -there are unitless prefactors multiplying this term in the +there are unitless prefactors multiplying this term in the equations. Note that these equations match what is in Dall'Amico (for θstar, @@ -50,9 +49,12 @@ function phase_change_source( τ::FT, ν::FT, θ_r::FT, - hydrology_cm::C, - earth_param_set::EP, -) where {FT, EP, C} + hydrology_earth_params::HEP, +) where {FT, HEP} + # Extract parameter sets from their container + hydrology_cm = hydrology_earth_params.hydrology_cm + earth_param_set = hydrology_earth_params.earth_param_set + _ρ_i = FT(LP.ρ_cloud_ice(earth_param_set)) _ρ_l = FT(LP.ρ_cloud_liq(earth_param_set)) _LH_f0 = FT(LP.LH_f0(earth_param_set)) @@ -71,6 +73,28 @@ function phase_change_source( return (θ_l - θstar) / τ end +""" + struct HydrologyEarthParameters{ + HCM <: AbstractSoilHydrologyClosure, + EP <: ClimaLand.Parameters.AbstractLandParameters, + } + +A wrapper type around the hydrology closure model and land parameter +structs. This is needed because of a type inference failure coming from +ClimaCore when multiple structs and fields are broadcasted over. +This struct circumvents that issue by wrapping the structs in a single +type, that can be unpacked within the broadcasted function. + +See github.com/CliMA/ClimaCore.jl/issues/2065 for more information +""" +struct HydrologyEarthParameters{ + HCM <: AbstractSoilHydrologyClosure, + EP <: ClimaLand.Parameters.AbstractLandParameters, +} + hydrology_cm::HCM + earth_param_set::EP +end +Base.broadcastable(x::HydrologyEarthParameters) = tuple(x) """ volumetric_heat_capacity( @@ -205,7 +229,7 @@ end ν::FT ) where {FT} -Compute the expression for relative saturation. +Compute the expression for relative saturation. This is referred to as θ_sat in Balland and Arp's paper. """ function relative_saturation(θ_l::FT, θ_i::FT, ν::FT) where {FT} diff --git a/test/standalone/Soil/soil_parameterizations.jl b/test/standalone/Soil/soil_parameterizations.jl index 096b8ee156..03094f1459 100644 --- a/test/standalone/Soil/soil_parameterizations.jl +++ b/test/standalone/Soil/soil_parameterizations.jl @@ -307,6 +307,10 @@ for FT in (Float32, Float64) vg_n = FT(1.4) hcm = vanGenuchten{FT}(; α = vg_α, n = vg_n) + # Wrap hydrology and earth parameters in one struct to avoid type inference failure + hydrology_earth_params = + ClimaLand.Soil.HydrologyEarthParameters(hcm, param_set) + K_sat = FT(1e-5) ν_ss_om = FT(0.1) ν_ss_gravel = FT(0.1) @@ -336,11 +340,26 @@ for FT in (Float32, Float64) ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set) τ = thermal_time.(ρc_s, Δz, parameters.κ_dry) @test ( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈ - (θ_l .- θ_star) ./ τ + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) ≈ (θ_l .- θ_star) ./ τ ) @test sum( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .> 0.0, + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) .> 0.0, ) == 3 # try θ_l = 0.1 @@ -354,8 +373,15 @@ for FT in (Float32, Float64) ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set) τ = thermal_time.(ρc_s, Δz, parameters.κ_dry) @test ( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈ - zeros(FT, 3) + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) ≈ zeros(FT, 3) ) @test (θ_star ≈ θ_l) @@ -370,11 +396,26 @@ for FT in (Float32, Float64) ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set) τ = thermal_time.(ρc_s, Δz, parameters.κ_dry) @test ( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈ - (θ_l .- θ_star) ./ τ + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) ≈ (θ_l .- θ_star) ./ τ ) @test sum( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .< 0.0, + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) .< 0.0, ) == 2 @@ -388,11 +429,26 @@ for FT in (Float32, Float64) ρc_s = volumetric_heat_capacity.(θ_l, θ_i, parameters.ρc_ds, param_set) τ = thermal_time.(ρc_s, Δz, parameters.κ_dry) @test ( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) ≈ - (θ_l .- θ_star) ./ τ + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) ≈ (θ_l .- θ_star) ./ τ ) @test sum( - phase_change_source.(θ_l, θ_i, T, τ, ν, θ_r, hcm, param_set) .> 0.0, + phase_change_source.( + θ_l, + θ_i, + T, + τ, + ν, + θ_r, + hydrology_earth_params, + ) .> 0.0, ) == 2 end end