Merge pull request #75 from TUM-PIK-ESM/bg/gpu-fixes

Fixes for type stability/inference issues affecting GPU support

Author: bgroenks96

src/abstract_variables.jl

@@ -393,11 +393,10 @@ Helper method that selects only closure (auxiliary) variables declared on `obj`.
 @inline closure_variables(obj) = closure_variables(variables(obj))
 @inline function closure_variables(vars::Tuple)
     progvars = prognostic_variables(vars)
-    closure_vars = mapreduce(var -> variables(var.closure), tuplejoin, progvars, init = ())
-    return deduplicate_vars(closure_vars)
+    all_closure_vars = fastmap(var -> variables(var.closure), progvars)
+    return deduplicate_vars(tuplejoin(all_closure_vars...))
 end
 
-
 function Base.NamedTuple(vars::Tuple{Vararg{Union{AbstractVariable, Namespace}}})
     keys = map(varname, vars)
     return NamedTuple{keys}(vars)

src/processes/soil/hydrology/soil_hydraulic_properties.jl

@@ -5,7 +5,7 @@
 
 Base type for unsaturated hydraulic conductivity parameterizations.
 """
-abstract type AbstractUnsatK end
+abstract type AbstractUnsatK{NF} end
 
 """
     get_swrc(::AbstractUnsatK)
@@ -63,7 +63,7 @@ measurements of hydraulic properites are available.
 Properties:
 $TYPEDFIELDS
 """
-@kwdef struct ConstantSoilHydraulics{NF, RC, UnsatK} <: AbstractSoilHydraulics{NF, RC, UnsatK}
+@kwdef struct ConstantSoilHydraulics{NF, RC, UnsatK <: AbstractUnsatK{NF}} <: AbstractSoilHydraulics{NF, RC, UnsatK}
     "Soil water retention curve"
     swrc::RC
 
@@ -78,12 +78,6 @@ $TYPEDFIELDS
 
     "Prescribed wilting point [-]"
     wilting_point::NF = 0.05
-
-    # TODO: Remove once FreezeCurves.jl allows for generic type bounds
-    function ConstantSoilHydraulics(swrc::SWRC, unsat_hydraulic_cond::AbstractUnsatK, args::NF...) where {NF}
-        adapted_swrc = adapt(NumberFormatAdaptor{NF}(), ustrip(swrc))
-        return new{NF, typeof(adapted_swrc), typeof(unsat_hydraulic_cond)}(adapted_swrc, unsat_hydraulic_cond, args...)
-    end
 end
 
 function ConstantSoilHydraulics(
@@ -92,7 +86,8 @@ function ConstantSoilHydraulics(
         unsat_hydraulic_cond = UnsatKLinear(NF),
         kwargs...
     ) where {NF}
-    return ConstantSoilHydraulics(; swrc, unsat_hydraulic_cond, kwargs...)
+    swrc = adapt(NumberFormatAdaptor{NF}(), ustrip(swrc))
+    return ConstantSoilHydraulics{NF, typeof(swrc), typeof(unsat_hydraulic_cond)}(; swrc, unsat_hydraulic_cond, kwargs...)
 end
 
 @inline saturated_hydraulic_conductivity(hydraulics::ConstantSoilHydraulics, args...) = hydraulics.sat_hydraulic_cond
@@ -110,7 +105,7 @@ and wilting point as a function of soil texture.
 Properties:
 $TYPEDFIELDS
 """
-@kwdef struct SoilHydraulicsSURFEX{NF, RC, UnsatK} <: AbstractSoilHydraulics{NF, RC, UnsatK}
+@kwdef struct SoilHydraulicsSURFEX{NF, RC, UnsatK <: AbstractUnsatK{NF}} <: AbstractSoilHydraulics{NF, RC, UnsatK}
     "Soil water retention curve"
     swrc::RC
 
@@ -128,12 +123,6 @@ $TYPEDFIELDS
 
     "Exponent of field capacity adjustment due to clay content [-]"
     field_capacity_exp::NF = 0.35
-
-    # TODO: Remove once FreezeCurves.jl allows for generic type bounds
-    function SoilHydraulicsSURFEX(swrc::SWRC, unsat_hydraulic_cond::AbstractUnsatK, args::NF...) where {NF}
-        adapted_swrc = adapt(NumberFormatAdaptor{NF}(), ustrip(swrc))
-        return new{NF, typeof(adapted_swrc), typeof(unsat_hydraulic_cond)}(adapted_swrc, unsat_hydraulic_cond, args...)
-    end
 end
 
 function SoilHydraulicsSURFEX(
@@ -142,7 +131,8 @@ function SoilHydraulicsSURFEX(
         unsat_hydraulic_cond = UnsatKLinear(NF),
         kwargs...
     ) where {NF}
-    return SoilHydraulicsSURFEX(; swrc, unsat_hydraulic_cond, kwargs...)
+    swrc = adapt(NumberFormatAdaptor{NF}(), ustrip(swrc))
+    return SoilHydraulicsSURFEX{NF, typeof(swrc), typeof(unsat_hydraulic_cond)}(; swrc, unsat_hydraulic_cond, kwargs...)
 end
 
 # TODO: this is not quite correct, SURFEX uses a hydraulic conductivity function that decreases exponentially with depth
@@ -169,7 +159,7 @@ end
 Simple formulation of hydraulic conductivity as a linear function of the liquid water saturated fraction,
 i.e. `soil.water / (soil.water + soil.ice + soil.air)`.
 """
-struct UnsatKLinear{NF} <: AbstractUnsatK end
+struct UnsatKLinear{NF} <: AbstractUnsatK{NF} end
 
 UnsatKLinear(::Type{NF}) where {NF} = UnsatKLinear{NF}()
 
@@ -194,7 +184,7 @@ volumetric fractions, assumed to include those of water, ice, and air.
 
 See van Genuchten (1980) and Westermann et al. (2023).
 """
-struct UnsatKVanGenuchten{NF} <: AbstractUnsatK
+struct UnsatKVanGenuchten{NF} <: AbstractUnsatK{NF}
     "Exponential scaling factor for ice impedance"
     impedance::NF
 end

src/processes/soil/stratigraphy/soil_volume.jl

@@ -88,7 +88,7 @@ $TYPEDFIELDS
     organic::NF = zero(eltype(texture))
 
     function MineralOrganic(texture::SoilTexture{NF}, organic::NF) where {NF}
-        @assert 0 <= organic <= 1 "organic content must be between zero and one"
+        @assert zero(NF) <= organic <= one(NF) "organic content must be between zero and one"
         return new{NF}(texture, organic)
     end
 end

src/state_variables.jl

@@ -211,11 +211,12 @@ end
 Retrieves all non-tendency `Field`s from `state` defined on the given `components`.
 """
 @inline function get_fields(state, components...; except = (;))
-    vars = mapreduce(tuplejoin, components, init = ()) do component
+    component_vars = fastmap(components) do component
         allvars = variables(component)
         closurevars = closure_variables(component)
         tuplejoin(allvars, closurevars)
     end
+    vars = tuplejoin(component_vars...)
     return ntdiff(get_fields(state, vars), except)
 end
 
@@ -225,7 +226,9 @@ end
 Retrieves all `Field`s from `state` corresponding to prognostic variables defined on the given `components`.
 """
 @inline function prognostic_fields(state, components...)
-    return get_fields(state, mapreduce(prognostic_variables, tuplejoin, components))
+    component_progvars = fastmap(prognostic_variables, components)
+    progvars = tuplejoin(component_progvars...)
+    return get_fields(state, progvars)
 end
 
 """
@@ -234,7 +237,9 @@ end
 Retrieves all `Field`s from `state` corresponding to tendencies defined on the given `components`.
 """
 @inline function tendency_fields(state, components...)
-    return get_fields(state.tendencies, mapreduce(prognostic_variables, tuplejoin, components))
+    component_progvars = fastmap(prognostic_variables, components)
+    progvars = tuplejoin(component_progvars...)
+    return get_fields(state.tendencies, progvars)
 end
 
 """
@@ -243,7 +248,9 @@ end
 Retrieves all `Field`s from `state` corresponding to auxiliary variables defined on the given `components`.
 """
 @inline function auxiliary_fields(state, components...)
-    return get_fields(state, mapreduce(auxiliary_variables, tuplejoin, components))
+    component_auxvars = fastmap(auxiliary_variables, components)
+    auxvars = tuplejoin(component_auxvars...)
+    return get_fields(state, auxvars)
 end
 
 """
@@ -252,7 +259,9 @@ end
 Retrieves all `Field`s from `state` corresponding to closure variables defined on the given `components`.
 """
 @inline function closure_fields(state, components...)
-    return get_fields(state, mapreduce(closure_variables, tuplejoin, components))
+    component_closurevars = fastmap(closure_variables, components)
+    closurevars = tuplejoin(component_closurevars...)
+    return get_fields(state, closurevars)
 end
 
 """
@@ -261,7 +270,9 @@ end
 Retrieves all `Field`s from `state` corresponding to input variables defined on the given `components`.
 """
 @inline function input_fields(state, components...)
-    return get_fields(state, mapreduce(input_variables, tuplejoin, components))
+    component_inputvars = fastmap(input_variables, components)
+    inputvars = tuplejoin(component_inputvars...)
+    return get_fields(state, inputvars)
 end
 
 # Initialization of StateVariables from models and processes

src/utils/utils.jl

@@ -22,7 +22,7 @@ convert_dt(Δt::Period) = Second(Δt).value
 
 Evaluates `x / (y + eps(NF))` if and only if `y != zero(y)`; returns `Inf` otherwise.
 """
-safediv(x::NF, y::NF) where {NF} = ifelse(iszero(y), Inf, x / (y + eps(NF)))
+safediv(x::NF, y::NF) where {NF} = ifelse(iszero(y), NF(Inf), x / (y + eps(NF)))
 
 # fastmap and fastiterate