add wrapper struct for type inference

This is required as a workaround to the type inference failure detailed
in CliMA/ClimaCore.jl#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

@@ -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
 -------

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}
 

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}

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