diff --git a/docs/src/API.md b/docs/src/API.md index 3cf36bd7c..8ef6768a6 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -388,6 +388,7 @@ Lightweight numerical utility functions shared across modules. ```@docs Utilities Utilities.clamp_to_nonneg +Utilities.promote_typeof Utilities.ϵ_numerics Utilities.ϵ_numerics_2M_M Utilities.ϵ_numerics_2M_N diff --git a/src/Microphysics2M.jl b/src/Microphysics2M.jl index 37faf5927..cb87aab4d 100644 --- a/src/Microphysics2M.jl +++ b/src/Microphysics2M.jl @@ -787,7 +787,9 @@ function rain_evaporation( (; pdf_r, evap)::CMP.SB2006, aps::CMP.AirProperties, tps::TDI.PS, q_tot, q_lcl, q_icl, q_rai, q_sno, ρ, N_rai, T, ) - FT = eltype(q_tot) + # the early return below must match the main path's type for any mix of + # plain-float and Dual arguments + FT = UT.promote_typeof(q_tot, q_lcl, q_icl, q_rai, q_sno, ρ, N_rai, T) ϵₘ = UT.ϵ_numerics_2M_M(FT) ϵₙ = UT.ϵ_numerics_2M_N(FT) diff --git a/src/Utilities.jl b/src/Utilities.jl index 87f8746a6..689f39202 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -9,6 +9,20 @@ module Utilities import UnrolledUtilities as UU export clamp_to_nonneg, ϵ_numerics, ϵ_numerics_2M_M, ϵ_numerics_2M_N, ϵ_numerics_P3_B +export promote_typeof + +""" + promote_typeof(args...) + +The common promoted type of the arguments' types. + +Use it to type early returns and sentinel values from all the arguments the +main-path result derives from. Typing them from a single argument +(`FT = eltype(q_tot)`-style) makes the function's return a union when a +caller mixes plain floats with `ForwardDiff.Dual`s (or float widths) across +arguments — non-concrete, heap-boxed, and silent. +""" +@inline promote_typeof(args...) = promote_type(map(typeof, args)...) export unrolled_logsumexp export sgs_weight_function, rime_mass_fraction, rime_density diff --git a/test/ad_compat_tests.jl b/test/ad_compat_tests.jl index 3a58b2e6d..e37b8b301 100644 --- a/test/ad_compat_tests.jl +++ b/test/ad_compat_tests.jl @@ -70,6 +70,22 @@ function test_ad_compatibility(FT) @test CM2.Γ_incl(FT(-0.25), FT(0.5)) isa FT end + @testset "rain_evaporation is concretely typed for mixed arguments ($FT)" begin + # the early return (no rain / supersaturated) must have the same type + # as the main path when species are Duals over a plain-float q_tot — + # a union here heap-boxes every call in the Jacobian hot loop + sb = mp.warm_rain.seifert_beheng + aps = mp.warm_rain.air_properties + z = D(0, 0) + # subsaturated (main path) and supersaturated (early return) states + for (q_tot, T) in ((FT(0.005), FT(288)), (FT(0.02), FT(288))) + t = @inferred CM2.rain_evaporation( + sb, aps, tps, q_tot, D(2e-4, 1), z, D(1e-4, 1), z, FT(1.05), D(4e4, 1), T, + ) + @test all(v -> v isa FD.Dual, values(t)) + end + end + @testset "BMT 2M+P3 Jacobian w.r.t. the 8 species ($FT)" begin function rhs(x, ρ, T, q_tot, logλ) t = BMT.bulk_microphysics_tendencies(